1 files changed, 143 insertions, 0 deletions
diff --git a/Simplify.hs b/Simplify.hs
new file mode 100644
index 0000000..7435464
--- /dev/null
+++ b/Simplify.hs
@@ -0,0 +1,143 @@
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE ViewPatterns #-}
+module Simplify where
+
+import Data.Monoid (Sum(..))
+
+import AST
+import ASTfunc
+import MonMap (MonMap)
+import qualified MonMap
+
+
+simplify :: Expr env t -> Expr env t
+simplify e = unSExpr (slr' . sli' . slr' $ toSExpr e)
+
+slr' :: SExpr env t -> SExpr env t
+slr' e = slr weakenId e
+
+-- simplify using let-rotate, and work away Unused bindings
+slr :: env :> env' -> SExpr env t -> SExpr env' t
+slr w (Auto expr) = case expr of
+    -- unused binding
+    Let lhs Unused e -> slr (weakenSkip lhs w) e
+
+    -- let-rotate
+    Let lhs (Auto (Let lhs2 rhs2 e2)) e
+      | Exists lhs' <- rebuildLHS lhs
+      -> let sh = sinkWithLHS lhs lhs' (weakenWithLHS lhs2)
+         in slr w (In (Let lhs2 rhs2 (In (Let lhs' e2 (applyShift sh e)))))
+
+    -- let-split
+    Let (L2 l1 l2) (Auto (Pair e1 e2)) e ->
+        slr w (In (Let l1 e1 (In (Let l2 (applyShift (weakenWithLHS l1) e2) e))))
+
+    -- redundant wildcard binding
+    Let (L0 _) _ e -> slr w e
+
+    Const t x -> In (Const t x)
+    Pair e1 e2 -> In (Pair (slr w e1) (slr w e2))
+    Nil -> In Nil
+    Prim op e -> In (Prim op (slr w e))
+    Var (V t i) -> In (Var (V t (w >:> i)))
+    Let lhs rhs e
+      | Exists lhs' <- rebuildLHS lhs
+      -> In (Let lhs' (slr w rhs) (slr (sinkWithLHS lhs lhs' w) e))
+slr _ Unused = error "slr: unused found outside Let RHS"
+
+-- simplify using linear inlining
+sli' :: SExpr env t -> SExpr env t
+sli' expr =
+    let (countmp, expr') = count 0 expr
+        countmp' = ((<= 1) . getSum) <$> countmp
+    in inline countmp' 0 (Inliner (In . Var)) expr'
+
+count :: Int -> SExpr env t -> (MonMap Int (Sum Int), SExpr env t)
+count d (Auto expr) = In <$> case expr of
+    Const t x -> return (Const t x)
+    Pair e1 e2 -> Pair <$> count d e1 <*> count d e2
+    Nil -> return Nil
+    Prim op e -> Prim op <$> count d e
+    Var (V t i) -> (MonMap.singleton (d - 1 - idxToInt i) 1, Var (V t i))
+    Let lhs rhs e -> Let lhs <$> count d rhs <*> count (d + lhsSize lhs) e
+count _ Unused = return Unused
+
+lhsSize :: LHS s t env env' -> Int
+lhsSize (L0 _) = 0
+lhsSize (L1 _) = 0
+lhsSize (L2 l1 l2) = lhsSize l1 + lhsSize l2
+
+data Inliner env = Inliner { unInliner :: forall t. V STy env t -> SExpr env t }
+
+inline :: MonMap Int Bool -> Int -> Inliner env -> SExpr env t -> SExpr env t
+inline countmp = go
+  where
+    go :: Int -> Inliner env -> SExpr env t -> SExpr env t
+    go d il (Auto expr) = case expr of
+        Const t x -> In (Const t x)
+        Pair e1 e2 -> In (Pair (go d il e1) (go d il e2))
+        Nil -> In Nil
+        Prim op e -> In (Prim op (go d il e))
+        Var v -> unInliner il v
+        Let (L1 t) rhs e
+          | Just True <- MonMap.lookup d countmp ->
+              let il' = Inliner
+                      (\case V _ Z -> applyShift weakenSucc rhs
+                             V t' (S i) -> applyShift weakenSucc (unInliner il (V t' i)))
+              in In (Let (L1 t) Unused (go (d + 1) il' e))
+        Let lhs rhs e ->
+            In (Let lhs (go d il rhs)
+                    (go (d + lhsSize lhs) (weakenILwithLHS lhs il) e))
+    go _ _ Unused = Unused
+
+weakenILwithLHS :: LHS STy t env env' -> Inliner env -> Inliner env'
+weakenILwithLHS (L0 _) il = il
+weakenILwithLHS (L1 _) il =
+    Inliner (\case v@(V _ Z) -> In (Var v)
+                   V t (S i) -> applyShift weakenSucc (unInliner il (V t i)))
+weakenILwithLHS (L2 l1 l2) il = weakenILwithLHS l2 (weakenILwithLHS l1 il)
+
+
+data SExpr env t where
+    Shift :: env :> env' -> PExpr SExpr env t -> SExpr env' t
+    In :: PExpr SExpr env t -> SExpr env t
+    Unused :: SExpr env t
+
+toSExpr :: Expr env t -> SExpr env t
+toSExpr (Expr e) = In (eInto toSExpr e)
+
+unSExpr :: SExpr env t -> Expr env t
+unSExpr (In e) = Expr (eInto unSExpr e)
+unSExpr (Shift w e) = Expr (shiftTraverse (eInto unSExpr) w e)
+unSExpr Unused = error "unSExpr: unused found"
+
+applyShift :: env :> env' -> SExpr env t -> SExpr env' t
+applyShift w (In e) = Shift w e
+applyShift w (Shift w' e) = Shift (w >. w') e
+applyShift _ Unused = Unused
+
+pattern Auto :: PExpr SExpr env t -> SExpr env t
+pattern Auto se <- (matchAuto -> Just se)
+{-# COMPLETE Auto, Unused #-}
+
+matchAuto :: SExpr env t -> Maybe (PExpr SExpr env t)
+matchAuto (In e) = Just e
+matchAuto (Shift w e) = Just (shiftTraverse id w e)
+matchAuto Unused = Nothing
+
+shiftTraverse :: (forall env' t'. PExpr SExpr env' t' -> PExpr expr env' t')
+              -> env :> env1 -> PExpr SExpr env t -> PExpr expr env1 t
+shiftTraverse f w e = case e of
+    Const t x -> Const t x
+    Pair e1 e2 -> f (Pair (applyShift w e1) (applyShift w e2))
+    Nil -> Nil
+    Prim op e1 -> f (Prim op (applyShift w e1))
+    Var v -> Var (weakenVar w v)
+    Let lhs rhs e1
+      | Exists lhs' <- rebuildLHS lhs
+      -> f (Let lhs' (applyShift w rhs)
+                     (applyShift (sinkWithLHS lhs lhs' w) e1))