simplifyFix

1 files changed, 50 insertions, 37 deletions

diff --git a/src/Simplify.hs b/src/Simplify.hs
index 5829a8b..cfbdbb9 100644
--- a/src/Simplify.hs
+++ b/src/Simplify.hs
@@ -9,6 +9,9 @@
 {-# LANGUAGE TypeOperators #-}
 module Simplify where
 
+import Data.Function (fix)
+import Data.Monoid (Any(..))
+
 import AST
 import AST.Count
 import Data
@@ -19,23 +22,30 @@ simplifyN 0 = id
 simplifyN n = simplifyN (n - 1) . simplify
 
 simplify :: forall env t. KnownEnv env => Ex env t -> Ex env t
-simplify = let ?accumInScope = checkAccumInScope @env knownEnv in simplify'
+simplify = let ?accumInScope = checkAccumInScope @env knownEnv in snd . simplify'
+
+simplifyFix :: forall env t. KnownEnv env => Ex env t -> Ex env t
+simplifyFix =
+  let ?accumInScope = checkAccumInScope @env knownEnv
+  in fix $ \loop e ->
+            let (Any act, e') = simplify' e
+            in if act then loop e' else e'
 
-simplify' :: (?accumInScope :: Bool) => Ex env t -> Ex env t
+simplify' :: (?accumInScope :: Bool) => Ex env t -> (Any, Ex env t)
 simplify' = \case
   -- inlining
   ELet _ rhs body
     | cheapExpr rhs
-    -> simplify' (subst1 rhs body)
+    -> acted $ simplify' (subst1 rhs body)
 
     | Occ lexOcc runOcc <- occCount IZ body
     , ((not ?accumInScope || not (hasAdds rhs)) && lexOcc <= One && runOcc <= One)  -- without effects, normal rules apply
           || (lexOcc == One && runOcc == One)  -- with effects, linear inlining is still allowed, but weakening is not
-    -> simplify' (subst1 rhs body)
+    -> acted $ simplify' (subst1 rhs body)
 
   -- let splitting
   ELet _ (EPair _ a b) body ->
-    simplify' $
+    acted $ simplify' $
       ELet ext a $
       ELet ext (weakenExpr WSink b) $
         subst (\_ t -> \case IZ -> EPair ext (EVar ext (typeOf a) (IS IZ)) (EVar ext (typeOf b) IZ)
@@ -44,10 +54,10 @@ simplify' = \case
 
   -- let rotation
   ELet _ (ELet _ rhs a) b ->
-    simplify' $
+    acted $ simplify' $
       ELet ext rhs $
       ELet ext a $
-        weakenExpr (WCopy WSink) (simplify' b)
+        weakenExpr (WCopy WSink) (snd (simplify' b))
 
   -- beta rules for products
   EFst _ (EPair _ e _) -> simplify' e
@@ -72,36 +82,39 @@ simplify' = \case
 
   -- TODO: accum of zero, plus of zero
 
-  EVar _ t i -> EVar ext t i
-  ELet _ a b -> ELet ext (simplify' a) (simplify' b)
-  EPair _ a b -> EPair ext (simplify' a) (simplify' b)
-  EFst _ e -> EFst ext (simplify' e)
-  ESnd _ e -> ESnd ext (simplify' e)
-  ENil _ -> ENil ext
-  EInl _ t e -> EInl ext t (simplify' e)
-  EInr _ t e -> EInr ext t (simplify' e)
-  ECase _ e a b -> ECase ext (simplify' e) (simplify' a) (simplify' b)
-  ENothing _ t -> ENothing ext t
-  EJust _ e -> EJust ext (simplify' e)
-  EMaybe _ a b e -> EMaybe ext (simplify' a) (simplify' b) (simplify' e)
-  EConstArr _ n t v -> EConstArr ext n t v
-  EBuild1 _ a b -> EBuild1 ext (simplify' a) (simplify' b)
-  EBuild _ n a b -> EBuild ext n (simplify' a) (simplify' b)
-  EFold1Inner _ a b c -> EFold1Inner ext (simplify' a) (simplify' b) (simplify' c)
-  ESum1Inner _ e -> ESum1Inner ext (simplify' e)
-  EUnit _ e -> EUnit ext (simplify' e)
-  EReplicate1Inner _ a b -> EReplicate1Inner ext (simplify' a) (simplify' b)
-  EConst _ t v -> EConst ext t v
-  EIdx0 _ e -> EIdx0 ext (simplify' e)
-  EIdx1 _ a b -> EIdx1 ext (simplify' a) (simplify' b)
-  EIdx _ a b -> EIdx ext (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)
-  EAccum i e1 e2 e3 -> EAccum i (simplify' e1) (simplify' e2) (simplify' e3)
-  EZero t -> EZero t
-  EPlus t a b -> EPlus t (simplify' a) (simplify' b)
-  EError t s -> EError t s
+  EVar _ t i -> pure $ EVar ext t i
+  ELet _ a b -> ELet ext <$> simplify' a <*> simplify' b
+  EPair _ a b -> EPair ext <$> simplify' a <*> simplify' b
+  EFst _ e -> EFst ext <$> simplify' e
+  ESnd _ e -> ESnd ext <$> simplify' e
+  ENil _ -> pure $ ENil ext
+  EInl _ t e -> EInl ext t <$> simplify' e
+  EInr _ t e -> EInr ext t <$> simplify' e
+  ECase _ e a b -> ECase ext <$> simplify' e <*> simplify' a <*> simplify' b
+  ENothing _ t -> pure $ ENothing ext t
+  EJust _ e -> EJust ext <$> simplify' e
+  EMaybe _ a b e -> EMaybe ext <$> simplify' a <*> simplify' b <*> simplify' e
+  EConstArr _ n t v -> pure $ EConstArr ext n t v
+  EBuild1 _ a b -> EBuild1 ext <$> simplify' a <*> simplify' b
+  EBuild _ n a b -> EBuild ext n <$> simplify' a <*> simplify' b
+  EFold1Inner _ a b c -> EFold1Inner ext <$> simplify' a <*> simplify' b <*> simplify' c
+  ESum1Inner _ e -> ESum1Inner ext <$> simplify' e
+  EUnit _ e -> EUnit ext <$> simplify' e
+  EReplicate1Inner _ a b -> EReplicate1Inner ext <$> simplify' a <*> simplify' b
+  EConst _ t v -> pure $ EConst ext t v
+  EIdx0 _ e -> EIdx0 ext <$> simplify' e
+  EIdx1 _ a b -> EIdx1 ext <$> simplify' a <*> simplify' b
+  EIdx _ a b -> EIdx ext <$> 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)
+  EAccum i e1 e2 e3 -> EAccum i <$> simplify' e1 <*> simplify' e2 <*> simplify' e3
+  EZero t -> pure $ EZero t
+  EPlus t a b -> EPlus t <$> simplify' a <*> simplify' b
+  EError t s -> pure $ EError t s
+
+acted :: (Any, a) -> (Any, a)
+acted (_, x) = (Any True, x)
 
 cheapExpr :: Expr x env t -> Bool
 cheapExpr = \case