Pretty print let bindings in do notation

3 files changed, 55 insertions, 37 deletions

diff --git a/src/AST/Pretty.hs b/src/AST/Pretty.hs
index 289294d..6bc75ed 100644
--- a/src/AST/Pretty.hs
+++ b/src/AST/Pretty.hs
@@ -63,26 +63,6 @@ ppExpr' :: Int -> SVal env -> Expr x env t -> M ShowS
 ppExpr' d val = \case
   EVar _ _ i -> return $ showString $ getConst $ valprj val i
 
-  etop@ELet{} -> 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 ((name, rhs') : binds, core)
-        collect val' e = ([],) <$> ppExpr' 0 val' e
-
-    (binds, core) <- collect val etop
-    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
-
   EPair _ a b -> do
     a' <- ppExpr' 0 val a
     b' <- ppExpr' 0 val b
@@ -193,31 +173,62 @@ ppExpr' d val = \case
     e' <- ppExpr' 11 val e
     return $ showParen (d > 10) $ showString ("return ") . e'
 
-  etop@(EMBind _ _) -> do
-    let collect :: SVal env -> Expr x env t -> M ([(String, ShowS)], ShowS)
-        collect val' (EMBind lhs cont) = do
-          let STEVM _ t = typeOf lhs
-          name <- genNameIfUsedIn t IZ cont
-          (binds, core) <- collect (VPush (Const name) val') cont
-          lhs' <- ppExpr' 0 val' lhs
-          return ((name, lhs') : binds, core)
-        collect val' e = ([],) <$> ppExpr' 0 val' e
-
-    (binds, core) <- collect val etop
-    return $ showParen (d > 0) $
-      showString "do { "
-      . foldr (.) id (intersperse (showString " ; ")
-                        (map (\(name, rhs) -> showString (name ++ " <- ") . rhs) binds))
-      . showString " ; " . core . showString " }"
+  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 IZ b
+  --   name <- genNameIfUsedIn t IZ b
   --   b' <- ppExpr' 0 (VPush (Const name) val) b
   --   return $ showParen (d > 10) $ a' . showString (" >>= \\" ++ name ++ " -> ") . b'
 
   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)
+      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)
+      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 " }"
+
 data Fixity = Prefix | Infix
   deriving (Show)
 
diff --git a/src/Example.hs b/src/Example.hs
index ee07edf..ee91981 100644
--- a/src/Example.hs
+++ b/src/Example.hs
@@ -92,6 +92,10 @@ ex4 =
 senv5 :: SList STy [TScal TF32, TEither (TScal TF32) (TScal TF32)]
 senv5 = STScal STF32 `SCons` STEither (STScal STF32) (STScal STF32) `SCons` SNil
 
+descr5 :: Storage a -> Storage b
+       -> Descr [TScal TF32, TEither (TScal TF32) (TScal TF32)] [b, a]
+descr5 a b = DTop `DPush` (STEither (STScal STF32) (STScal STF32), a) `DPush` (STScal STF32, b)
+
 -- x:R+R y:R |- case x of {inl a -> a * y ; inr b -> b * (y + 1)}
 ex5 :: Ex [TScal TF32, TEither (TScal TF32) (TScal TF32)] (TScal TF32)
 ex5 =
diff --git a/src/Simplify.hs b/src/Simplify.hs
index 16a3e1d..cbeee75 100644
--- a/src/Simplify.hs
+++ b/src/Simplify.hs
@@ -54,6 +54,9 @@ simplify = \case
   -- bind-let commute
   EMBind (ELet _ a b) c -> simplify (ELet ext a (EMBind b (weakenExpr (WCopy WSink) c)))
 
+  -- return-let commute
+  EMReturn env (ELet _ a b) -> simplify (ELet ext a (EMReturn env 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)