Generate EOneHot in D[EIdx]

This generates a one-hot for the zero-dimensional inner array because
indexing one level further to the actual element is too difficult. But
this should simplify away fine.

1 files changed, 21 insertions, 0 deletions

diff --git a/src/AST.hs b/src/AST.hs
index b9b10ad..60fc5ad 100644
--- a/src/AST.hs
+++ b/src/AST.hs
@@ -9,8 +9,10 @@
 {-# LANGUAGE PolyKinds #-}
 {-# LANGUAGE QuantifiedConstraints #-}
 {-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE StandaloneDeriving #-}
 {-# LANGUAGE StandaloneKindSignatures #-}
+{-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE TypeOperators #-}
 {-# LANGUAGE UndecidableInstances #-}
@@ -379,3 +381,22 @@ ezip a b =
                              (EVar ext (tTup (sreplicate n tIx)) IZ))
                    (EIdx ext (EVar ext (STArr n t2) (IS IZ))
                              (EVar ext (tTup (sreplicate n tIx)) IZ))
+
+arrIdxToAcIdx :: proxy t -> SNat n -> Ex env (Tup (Replicate n TIx)) -> Ex env (AcIdx (TArr n t) n)
+arrIdxToAcIdx = \p (n :: SNat n) e -> case lemPlusZero @n of Refl -> go p n SZ e (ENil ext)
+  where
+    -- symbolic version of 'invert' in Interpreter
+    go :: forall n m t env proxy. proxy t -> SNat n -> SNat m
+       -> Ex env (Tup (Replicate n TIx)) -> Ex env (AcIdx (TArr m t) m) -> Ex env (AcIdx (TArr (n + m) t) (n + m))
+    go _ SZ _ _ acidx = acidx
+    go p (SS n) m idx acidx
+      | Refl <- lemPlusSuccRight @n @m
+      = ELet ext idx $
+          go p n (SS m)
+             (EFst ext (EVar ext (typeOf idx) IZ))
+             (EPair ext (ESnd ext (EVar ext (typeOf idx) IZ))
+                        (weakenExpr WSink acidx))
+
+lemAcValArrN :: proxy t -> SNat n -> AcValArr n t n :~: TArr Z t
+lemAcValArrN _ SZ = Refl
+lemAcValArrN p (SS n) | Refl <- lemAcValArrN p n = Refl