1 files changed, 10 insertions, 48 deletions

diff --git a/src/Simplify.hs b/src/Simplify.hs
index cb649d5..acc2392 100644
--- a/src/Simplify.hs
+++ b/src/Simplify.hs
@@ -6,8 +6,13 @@
 module Simplify where
 
 import AST
+import AST.Count
 
 
+simplifyN :: Int -> Ex env t -> Ex env t
+simplifyN 0 = id
+simplifyN n = simplifyN (n - 1) . simplify
+
 simplify :: Ex env t -> Ex env t
 simplify = \case
   -- inlining
@@ -28,9 +33,11 @@ simplify = \case
                              IS i -> EVar ext t (IS (IS i)))
               body
 
+  -- beta rules for products
   EFst _ (EPair _ e _) -> simplify e
   ESnd _ (EPair _ _ e) -> simplify e
 
+  -- beta rules for coproducts
   ECase _ (EInl _ _ e) rhs _ -> simplify (ELet ext e rhs)
   ECase _ (EInr _ _ e) _ rhs -> simplify (ELet ext e rhs)
 
@@ -38,6 +45,9 @@ simplify = \case
 
   -- TODO: constant folding for operations
 
+  -- eta rule for return+bind
+  EMBind (EMReturn _ a) b -> simplify (ELet ext a b)
+
   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)
@@ -67,54 +77,6 @@ cheapExpr = \case
   EConst{} -> True
   _ -> False
 
-data Count = Zero | One | Many
-  deriving (Show, Eq, Ord)
-
-instance Semigroup Count where
-  Zero <> n = n
-  n <> Zero = n
-  _ <> _ = Many
-instance Monoid Count where
-  mempty = Zero
-
-data Occ = Occ { _occLexical :: Count
-               , _occRuntime :: Count }
-instance Semigroup Occ where Occ a b <> Occ c d = Occ (a <> c) (b <> d)
-instance Monoid Occ    where mempty = Occ mempty mempty
-
--- | One of the two branches is taken
-(<||>) :: Occ -> Occ -> Occ
-Occ l1 r1 <||> Occ l2 r2 = Occ (l1 <> l2) (max r1 r2)
-
--- | This code is executed many times
-scaleMany :: Occ -> Occ
-scaleMany (Occ l _) = Occ l Many
-
-occCount :: Idx env a -> Expr x env t -> Occ
-occCount idx = \case
-  EVar _ _ i | idx2int i == idx2int idx -> Occ One One
-             | otherwise -> mempty
-  ELet _ rhs body -> occCount idx rhs <> occCount (IS idx) body
-  EPair _ a b -> occCount idx a <> occCount idx b
-  EFst _ e -> occCount idx e
-  ESnd _ e -> occCount idx e
-  ENil _ -> mempty
-  EInl _ _ e -> occCount idx e
-  EInr _ _ e -> occCount idx e
-  ECase _ e a b -> occCount idx e <> (occCount (IS idx) a <||> occCount (IS idx) b)
-  EBuild1 _ a b -> occCount idx a <> scaleMany (occCount (IS idx) b)
-  EBuild _ es e -> foldMap (occCount idx) es <> scaleMany (occCount (wsinkN (vecLength es) @> idx) e)
-  EFold1 _ a b -> scaleMany (occCount (IS (IS idx)) a) <> occCount idx b
-  EConst{} -> mempty
-  EIdx1 _ a b -> occCount idx a <> occCount idx b
-  EIdx _ e es -> occCount idx e <> foldMap (occCount idx) es
-  EOp _ _ e -> occCount idx e
-  EMOne _ _ e -> occCount idx e
-  EMScope e -> occCount idx e
-  EMReturn _ e -> occCount idx e
-  EMBind a b -> occCount idx a <> occCount (IS idx) b
-  EError{} -> mempty
-
 subst1 :: Expr x env a -> Expr x (a : env) t -> Expr x env t
 subst1 repl = subst $ \x t -> \case IZ -> repl
                                     IS i -> EVar x t i