Migrate to accumulators (mostly removing EVM code)

3 files changed, 63 insertions, 82 deletions

diff --git a/src/AST/Count.hs b/src/AST/Count.hs
index de04b5f..f66b809 100644
--- a/src/AST/Count.hs
+++ b/src/AST/Count.hs
@@ -1,7 +1,12 @@
-{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE DerivingVia #-}
 {-# LANGUAGE GADTs #-}
+{-# LANGUAGE LambdaCase #-}
 module AST.Count where
 
+import GHC.Generics (Generic, Generically(..))
+
 import AST
 import Data
 
@@ -18,9 +23,8 @@ instance Monoid Count where
 
 data Occ = Occ { _occLexical :: Count
                , _occRuntime :: Count }
-  deriving (Eq)
-instance Semigroup Occ where Occ a b <> Occ c d = Occ (a <> c) (b <> d)
-instance Monoid Occ    where mempty = Occ mempty mempty
+  deriving (Eq, Generic)
+  deriving (Semigroup, Monoid) via Generically Occ
 
 -- | One of the two branches is taken
 (<||>) :: Occ -> Occ -> Occ
@@ -49,8 +53,6 @@ occCount idx = \case
   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
+  EWith a b -> occCount idx a <> occCount (IS idx) b
+  EAccum a b e -> occCount idx a <> occCount idx b <> occCount idx e
   EError{} -> mempty
diff --git a/src/AST/Pretty.hs b/src/AST/Pretty.hs
index 1ffa980..3473131 100644
--- a/src/AST/Pretty.hs
+++ b/src/AST/Pretty.hs
@@ -41,14 +41,20 @@ instance Monad M where { M f >>= g = M (\i -> let (x, j) = f i in runM (g x) j)
 genId :: M Int
 genId = M (\i -> (i, i + 1))
 
+genName' :: String -> M String
+genName' prefix = (prefix ++) . show <$> genId
+
 genName :: M String
-genName = ('x' :) . show <$> genId
+genName = genName' "x"
 
-genNameIfUsedIn :: STy a -> Idx env a -> Expr x env t -> M String
-genNameIfUsedIn ty idx ex
+genNameIfUsedIn' :: String -> STy a -> Idx env a -> Expr x env t -> M String
+genNameIfUsedIn' prefix ty idx ex
   | occCount idx ex == mempty = case ty of STNil -> return "()"
                                            _ -> return "_"
-  | otherwise                 = genName
+  | otherwise                 = genName' prefix
+
+genNameIfUsedIn :: STy a -> Idx env a -> Expr x env t -> M String
+genNameIfUsedIn = genNameIfUsedIn' "x"
 
 ppExpr :: SList STy env -> Expr x env t -> String
 ppExpr senv e = fst (runM (mkVal senv >>= \val -> ppExpr' 0 val e) 1) ""
@@ -64,6 +70,8 @@ ppExpr' :: Int -> SVal env -> Expr x env t -> M ShowS
 ppExpr' d val = \case
   EVar _ _ i -> return $ showString $ getConst $ valprj val i
 
+  e@ELet{} -> ppExprLet d val e
+
   EPair _ a b -> do
     a' <- ppExpr' 0 val a
     b' <- ppExpr' 0 val b
@@ -155,80 +163,46 @@ ppExpr' d val = \case
                 (Prefix, s) -> s
     return $ showParen (d > 10) $ showString (ops ++ " ") . e'
 
-  EMOne venv i e -> do
-    let venvlen = length (unSList venv)
-        varname = 'v' : show (venvlen - idx2int i)
-    e' <- ppExpr' 11 val e
+  EWith e1 e2 -> do
+    e1' <- ppExpr' 11 val e1
+    let STArr n t = typeOf e1
+    name <- genNameIfUsedIn' "ac" (STAccum n t) IZ e2
+    e2' <- ppExpr' 11 (VPush (Const name) val) e2
     return $ showParen (d > 10) $
-      showString ("one " ++ show varname ++ " ") . e'
+      showString "with " . e1' . showString (" (\\" ++ name ++ " -> ")
+      . e2' . showString ")"
 
-  EMScope e -> do
-    let venv = case typeOf e of STEVM v _ -> v
-        venvlen = length (unSList venv)
-        varname = 'v' : show venvlen
-    e' <- ppExpr' 11 val e
+  EAccum e1 e2 e3 -> do
+    e1' <- ppExpr' 11 val e1
+    e2' <- ppExpr' 11 val e2
+    e3' <- ppExpr' 11 val e3
     return $ showParen (d > 10) $
-      showString ("scope " ++ show varname ++ " ") . e'
-
-  EMReturn _ e -> do
-    e' <- ppExpr' 11 val e
-    return $ showParen (d > 10) $ showString ("return ") . e'
-
-  e@EMBind{} -> ppExprDo d val e
-  e@ELet{} -> ppExprDo d val e
-
-  -- EMBind a b -> do
-  --   let STEVM _ t = typeOf a
-  --   a' <- ppExpr' 0 val a
-  --   name <- genNameIfUsedIn t IZ b
-  --   b' <- ppExpr' 0 (VPush (Const name) val) b
-  --   return $ showParen (d > 10) $ a' . showString (" >>= \\" ++ name ++ " -> ") . b'
+      showString "accum " . e1' . showString " " . e2' . showString " " . e3'
 
   EError _ s -> return $ showParen (d > 10) $ showString ("error " ++ show s)
 
-data Binding = MonadBind String ShowS
-             | LetBind   String ShowS
-
-ppExprDo :: Int -> SVal env -> Expr x env t -> M ShowS
-ppExprDo d val etop = do
-  let collect :: SVal env -> Expr x env t -> M ([Binding], ShowS)
-      collect val' (EMBind lhs body) = do
-        let STEVM _ t = typeOf lhs
-        name <- genNameIfUsedIn t IZ body
-        (binds, core) <- collect (VPush (Const name) val') body
-        lhs' <- ppExpr' 0 val' lhs
-        return (MonadBind name lhs' : binds, core)
+ppExprLet :: Int -> SVal env -> Expr x env t -> M ShowS
+ppExprLet d val etop = do
+  let collect :: SVal env -> Expr x env t -> M ([(String, ShowS)], ShowS)
       collect val' (ELet _ rhs body) = do
         name <- genNameIfUsedIn (typeOf rhs) IZ body
-        (binds, core) <- collect (VPush (Const name) val') body
         rhs' <- ppExpr' 0 val' rhs
-        return (LetBind name rhs' : binds, core)
+        (binds, core) <- collect (VPush (Const name) val') body
+        return ((name, rhs') : binds, core)
       collect val' e = ([],) <$> ppExpr' 0 val' e
 
-      fromLet = \case LetBind n s -> Just (n, s) ; _ -> Nothing
-
   (binds, core) <- collect val etop
 
-  return $ showParen (d > 0) $ case traverse fromLet binds of
-    Just lbinds ->
-      let (open, close) = case lbinds of
-            [_] -> ("{ ", " }")
-            _ -> ("", "")
-      in showString ("let " ++ open)
-         . foldr (.) id
-             (intersperse (showString " ; ")
-                (map (\(name, rhs) -> showString (name ++ " = ") . rhs) lbinds))
-         . showString (close ++ " in ")
-         . core
-    Nothing ->
-      showString "do { "
-      . foldr (.) id
-          (intersperse (showString " ; ")
-             (map (\case MonadBind name rhs -> showString (name ++ " <- ") . rhs
-                         LetBind name rhs -> showString ("let { " ++ name ++ " = ") . rhs
-                                             . showString " }")
-                  binds))
-      . showString " ; " . core . showString " }"
+  let (open, close) = case binds of
+        [_] -> ("{ ", " }")
+        _ -> ("", "")
+  return $ showParen (d > 0) $
+    showString ("let " ++ open)
+    . foldr (.) id
+        (intersperse (showString " ; ")
+           (map (\(name, rhs) -> showString (name ++ " = ") . rhs) binds))
+    . showString (close ++ " in ")
+    . core
 
 data Fixity = Prefix | Infix
   deriving (Show)
diff --git a/src/AST/Weaken.hs b/src/AST/Weaken.hs
index 4b3016d..d992404 100644
--- a/src/AST/Weaken.hs
+++ b/src/AST/Weaken.hs
@@ -36,7 +36,7 @@ data env :> env' where
   WIdx :: Idx env t -> (t : env) :> env
 deriving instance Show (env :> env')
 
-infixr @>
+infixr 2 @>
 (@>) :: env :> env' -> Idx env t -> Idx env' t
 WId @> i = i
 WSink @> i = IS i
@@ -48,6 +48,7 @@ WClosed _ @> i = case i of {}
 WIdx j @> IZ = j
 WIdx _ @> IS i = i
 
+infixr 3 .>
 (.>) :: env2 :> env3 -> env1 :> env2 -> env1 :> env3
 (.>) = flip WThen
 
@@ -70,14 +71,18 @@ wCopies :: SList f bs -> env1 :> env2 -> Append bs env1 :> Append bs env2
 wCopies SNil w = w
 wCopies (SCons _ spine) w = WCopy (wCopies spine w)
 
--- wStack :: forall env b1 b2. b1 :> b2 -> Append b1 env :> Append b2 env
--- wStack WId = WId
--- wStack WSink = WSink
--- wStack (WCopy w) = WCopy (wStack @env w)
--- wStack (WPop w) = WPop (wStack @env w)
--- wStack (WThen w1 w2) = WThen (wStack @env w1) (wStack @env w2)
--- wStack (WClosed s) = wSinks s
--- wStack (WIdx i) = WIdx (_ i)
+wStack :: forall env b1 b2. b1 :> b2 -> Append b1 env :> Append b2 env
+wStack WId = WId
+wStack WSink = WSink
+wStack (WCopy w) = WCopy (wStack @env w)
+wStack (WPop w) = WPop (wStack @env w)
+wStack (WThen w1 w2) = WThen (wStack @env w1) (wStack @env w2)
+wStack (WClosed s) = wSinks s
+wStack (WIdx i) = WIdx (goIdx i)
+  where
+    goIdx :: Idx b t -> Idx (Append b env) t
+    goIdx IZ = IZ
+    goIdx (IS i') = IS (goIdx i')
 
 wRaiseAbove :: SList f env1 -> SList g env -> env1 :> Append env1 env
 wRaiseAbove SNil env = WClosed (slistMap (\_ -> Const ()) env)