1 files changed, 48 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)