WIP EOneHot

1 files changed, 3 insertions, 0 deletions

diff --git a/src/AST.hs b/src/AST.hs
index 71001e7..328a670 100644
--- a/src/AST.hs
+++ b/src/AST.hs
@@ -103,6 +103,7 @@ data Expr x env t where
   -- monoidal operations (to be desugared to regular operations after simplification)
   EZero :: STy t -> Expr x env (D2 t)
   EPlus :: STy t -> Expr x env (D2 t) -> Expr x env (D2 t) -> Expr x env (D2 t)
+  EOneHot :: STy t -> SNat i -> Expr x env (AcIdx t i) -> Expr x env (D2 (AcVal t i)) -> Expr x env (D2 t)
 
   -- partiality
   EError :: STy a -> String -> Expr x env a
@@ -206,6 +207,7 @@ typeOf = \case
 
   EZero t -> d2 t
   EPlus t _ _ -> d2 t
+  EOneHot t _ _ _ -> d2 t
 
   EError t _ -> t
 
@@ -277,6 +279,7 @@ subst' f w = \case
   EAccum i e1 e2 e3 -> EAccum i (subst' f w e1) (subst' f w e2) (subst' f w e3)
   EZero t -> EZero t
   EPlus t a b -> EPlus t (subst' f w a) (subst' f w b)
+  EOneHot t i a b -> EOneHot t i (subst' f w a) (subst' f w b)
   EError t s -> EError t s
   where
     sinkF :: (forall a. x a -> STy a -> (env' :> env2) -> Idx env a -> Expr x env2 a)