Add ext field to remaining AST constructors

1 files changed, 18 insertions, 18 deletions

diff --git a/src/AST.hs b/src/AST.hs
index 333f306..bcbb19a 100644
--- a/src/AST.hs
+++ b/src/AST.hs
@@ -111,17 +111,17 @@ data Expr x env t where
           -> Expr x env t
 
   -- accumulation effect
-  EWith :: Expr x env t -> Expr x (TAccum t : env) a -> Expr x env (TPair a t)
-  EAccum :: SNat i -> Expr x env (AcIdx t i) -> Expr x env (AcVal t i) -> Expr x env (TAccum t) -> Expr x env TNil
+  EWith :: x (TPair a t) -> Expr x env t -> Expr x (TAccum t : env) a -> Expr x env (TPair a t)
+  EAccum :: x TNil -> SNat i -> Expr x env (AcIdx t i) -> Expr x env (AcVal t i) -> Expr x env (TAccum t) -> Expr x env TNil
   -- EAccum1 :: Expr x env TIx -> Expr x env t -> Expr x env (TAccum (S Z) t) -> Expr x env TNil
 
   -- 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 (D2 t) i) -> Expr x env (AcVal (D2 t) i) -> Expr x env (D2 t)
+  EZero :: x (D2 t) -> STy t -> Expr x env (D2 t)
+  EPlus :: x (D2 t) -> STy t -> Expr x env (D2 t) -> Expr x env (D2 t) -> Expr x env (D2 t)
+  EOneHot :: x (D2 t) -> STy t -> SNat i -> Expr x env (AcIdx (D2 t) i) -> Expr x env (AcVal (D2 t) i) -> Expr x env (D2 t)
 
   -- partiality
-  EError :: STy a -> String -> Expr x env a
+  EError :: x a -> STy a -> String -> Expr x env a
 deriving instance (forall ty. Show (x ty)) => Show (Expr x env t)
 
 type Ex = Expr (Const ())
@@ -247,14 +247,14 @@ typeOf = \case
 
   ECustom _ _ _ _ e _ _ _ _ -> typeOf e
 
-  EWith e1 e2 -> STPair (typeOf e2) (typeOf e1)
-  EAccum _ _ _ _ -> STNil
+  EWith _ e1 e2 -> STPair (typeOf e2) (typeOf e1)
+  EAccum _ _ _ _ _ -> STNil
 
-  EZero t -> d2 t
-  EPlus t _ _ -> d2 t
-  EOneHot t _ _ _ -> d2 t
+  EZero _ t -> d2 t
+  EPlus _ t _ _ -> d2 t
+  EOneHot _ t _ _ _ -> d2 t
 
-  EError t _ -> t
+  EError _ t _ -> t
 
 -- unSNat :: SNat n -> Nat
 -- unSNat SZ = Z
@@ -322,12 +322,12 @@ subst' f w = \case
   EShape x e -> EShape x (subst' f w e)
   EOp x op e -> EOp x op (subst' f w e)
   ECustom x s t p a b c e1 e2 -> ECustom x s t p a b c (subst' f w e1) (subst' f w e2)
-  EWith e1 e2 -> EWith (subst' f w e1) (subst' (sinkF f) (WCopy w) e2)
-  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
+  EWith x e1 e2 -> EWith x (subst' f w e1) (subst' (sinkF f) (WCopy w) e2)
+  EAccum x i e1 e2 e3 -> EAccum x i (subst' f w e1) (subst' f w e2) (subst' f w e3)
+  EZero x t -> EZero x t
+  EPlus x t a b -> EPlus x t (subst' f w a) (subst' f w b)
+  EOneHot x t i a b -> EOneHot x t i (subst' f w a) (subst' f w b)
+  EError x t s -> EError x t s
   where
     sinkF :: (forall a. x a -> STy a -> (env' :> env2) -> Idx env a -> Expr x env2 a)
           -> x t -> STy t -> ((b : env') :> env2) -> Idx (b : env) t -> Expr x env2 t