Commutativity marker on fold1i

1 files changed, 7 insertions, 4 deletions

diff --git a/src/AST.hs b/src/AST.hs
index ecd4647..a4898c0 100644
--- a/src/AST.hs
+++ b/src/AST.hs
@@ -60,7 +60,7 @@ data Expr x env t where
   -- array operations
   EConstArr :: Show (ScalRep t) => x (TArr n (TScal t)) -> SNat n -> SScalTy t -> Array n (ScalRep t) -> Expr x env (TArr n (TScal t))
   EBuild :: x (TArr n t) -> SNat n -> Expr x env (Tup (Replicate n TIx)) -> Expr x (Tup (Replicate n TIx) : env) t -> Expr x env (TArr n t)
-  EFold1Inner :: x (TArr n t) -> Expr x (t : t : env) t -> Expr x env t -> Expr x env (TArr (S n) t) -> Expr x env (TArr n t)
+  EFold1Inner :: x (TArr n t) -> Commutative -> Expr x (t : t : env) t -> Expr x env t -> Expr x env (TArr (S n) t) -> Expr x env (TArr n t)
   ESum1Inner :: ScalIsNumeric t ~ True => x (TArr n (TScal t)) -> Expr x env (TArr (S n) (TScal t)) -> Expr x env (TArr n (TScal t))
   EUnit :: x (TArr Z t) -> Expr x env t -> Expr x env (TArr Z t)
   EReplicate1Inner :: x (TArr (S n) t) -> Expr x env TIx -> Expr x env (TArr n t) -> Expr x env (TArr (S n) t)
@@ -106,6 +106,9 @@ type Ex = Expr (Const ())
 ext :: Const () a
 ext = Const ()
 
+data Commutative = Commut | Noncommut
+  deriving (Show)
+
 type SOp :: Ty -> Ty -> Type
 data SOp a t where
   OAdd :: ScalIsNumeric a ~ True => SScalTy a -> SOp (TPair (TScal a) (TScal a)) (TScal a)
@@ -182,7 +185,7 @@ typeOf = \case
 
   EConstArr _ n t _ -> STArr n (STScal t)
   EBuild _ n _ e -> STArr n (typeOf e)
-  EFold1Inner _ _ _ e | STArr (SS n) t <- typeOf e -> STArr n t
+  EFold1Inner _ _ _ _ e | STArr (SS n) t <- typeOf e -> STArr n t
   ESum1Inner _ e | STArr (SS n) t <- typeOf e -> STArr n t
   EUnit _ e -> STArr SZ (typeOf e)
   EReplicate1Inner _ _ e | STArr n t <- typeOf e -> STArr (SS n) t
@@ -223,7 +226,7 @@ extOf = \case
   EMaybe x _ _ _ -> x
   EConstArr x _ _ _ -> x
   EBuild x _ _ _ -> x
-  EFold1Inner x _ _ _ -> x
+  EFold1Inner x _ _ _ _ -> x
   ESum1Inner x _ -> x
   EUnit x _ -> x
   EReplicate1Inner x _ _ -> x
@@ -292,7 +295,7 @@ subst' f w = \case
   EMaybe x a b e -> EMaybe x (subst' f w a) (subst' (sinkF f) (WCopy w) b) (subst' f w e)
   EConstArr x n t a -> EConstArr x n t a
   EBuild x n a b -> EBuild x n (subst' f w a) (subst' (sinkF f) (WCopy w) b)
-  EFold1Inner x a b c -> EFold1Inner x (subst' (sinkF (sinkF f)) (WCopy (WCopy w)) a) (subst' f w b) (subst' f w c)
+  EFold1Inner x cm a b c -> EFold1Inner x cm (subst' (sinkF (sinkF f)) (WCopy (WCopy w)) a) (subst' f w b) (subst' f w c)
   ESum1Inner x e -> ESum1Inner x (subst' f w e)
   EUnit x e -> EUnit x (subst' f w e)
   EReplicate1Inner x a b -> EReplicate1Inner x (subst' f w a) (subst' f w b)