Generic accumulators

1 files changed, 19 insertions, 11 deletions

diff --git a/src/Simplify.hs b/src/Simplify.hs
index 698c667..62a3a9c 100644
--- a/src/Simplify.hs
+++ b/src/Simplify.hs
@@ -8,8 +8,6 @@
 {-# LANGUAGE TypeOperators #-}
 module Simplify where
 
-import Data.Monoid
-
 import AST
 import AST.Count
 import Data
@@ -45,9 +43,10 @@ simplify' = \case
 
   -- let rotation
   ELet _ (ELet _ rhs a) b ->
-    ELet ext (simplify' rhs) $
-    ELet ext (simplify' a) $
-      weakenExpr (WCopy WSink) (simplify' b)
+    simplify' $
+      ELet ext rhs $
+      ELet ext a $
+        weakenExpr (WCopy WSink) (simplify' b)
 
   -- beta rules for products
   EFst _ (EPair _ e _) -> simplify' e
@@ -57,6 +56,13 @@ simplify' = \case
   ECase _ (EInl _ _ e) rhs _ -> simplify' (ELet ext e rhs)
   ECase _ (EInr _ _ e) _ rhs -> simplify' (ELet ext e rhs)
 
+  -- let floating to facilitate beta reduction
+  EFst _ (ELet _ rhs body) -> simplify' (ELet ext rhs (EFst ext body))
+  ESnd _ (ELet _ rhs body) -> simplify' (ELet ext rhs (ESnd ext body))
+  ECase _ (ELet _ rhs body) e1 e2 -> simplify' (ELet ext rhs (ECase ext body (weakenExpr (WCopy WSink) e1) (weakenExpr (WCopy WSink) e2)))
+  EIdx0 _ (ELet _ rhs body) -> simplify' (ELet ext rhs (EIdx0 ext body))
+  EIdx1 _ (ELet _ rhs body) e -> simplify' (ELet ext rhs (EIdx1 ext body (weakenExpr WSink e)))
+
   -- TODO: array indexing (index of build, index of fold)
 
   -- TODO: constant folding for operations
@@ -74,14 +80,15 @@ simplify' = \case
   EBuild _ n a b -> EBuild ext n (simplify' a) (simplify' b)
   EFold1 _ a b -> EFold1 ext (simplify' a) (simplify' b)
   EUnit _ e -> EUnit ext (simplify' e)
-  EReplicate _ e -> EReplicate ext (simplify' e)
+  -- EReplicate _ e -> EReplicate ext (simplify' e)
   EConst _ t v -> EConst ext t v
   EIdx0 _ e -> EIdx0 ext (simplify' e)
   EIdx1 _ a b -> EIdx1 ext (simplify' a) (simplify' b)
-  EIdx _ e es -> EIdx ext (simplify' e) (fmap simplify' es)
+  EIdx _ n a b -> EIdx ext n (simplify' a) (simplify' b)
+  EShape _ e -> EShape ext (simplify' e)
   EOp _ op e -> EOp ext op (simplify' e)
   EWith e1 e2 -> EWith (simplify' e1) (let ?accumInScope = True in simplify' e2)
-  EAccum1 e1 e2 e3 -> EAccum1 (simplify' e1) (simplify' e2) (simplify' e3)
+  EAccum i e1 e2 e3 -> EAccum i (simplify' e1) (simplify' e2) (simplify' e3)
   EError t s -> EError t s
 
 cheapExpr :: Expr x env t -> Bool
@@ -108,14 +115,15 @@ hasAdds = \case
   EBuild _ _ a b -> hasAdds a || hasAdds b
   EFold1 _ a b -> hasAdds a || hasAdds b
   EUnit _ e -> hasAdds e
-  EReplicate _ e -> hasAdds e
+  -- EReplicate _ e -> hasAdds e
   EConst _ _ _ -> False
   EIdx0 _ e -> hasAdds e
   EIdx1 _ a b -> hasAdds a || hasAdds b
-  EIdx _ e es -> hasAdds e || getAny (foldMap (Any . hasAdds) es)
+  EIdx _ _ a b -> hasAdds a || hasAdds b
+  EShape _ e -> hasAdds e
   EOp _ _ e -> hasAdds e
   EWith a b -> hasAdds a || hasAdds b
-  EAccum1 _ _ _ -> True
+  EAccum _ _ _ _ -> True
   EError _ _ -> False
 
 checkAccumInScope :: SList STy env -> Bool