Fix pretty-printer

1 files changed, 45 insertions, 18 deletions

diff --git a/src/SPL/PrettyPrinter.hs b/src/SPL/PrettyPrinter.hs
index 48cdf9d..e9a3300 100644
--- a/src/SPL/PrettyPrinter.hs
+++ b/src/SPL/PrettyPrinter.hs
@@ -15,31 +15,46 @@ prettyPrint t = result $ execState (pprint t) (St 0 "")
 
 -- For Extensibility
 data St = St
-  { indent :: Int
-  , result :: String
+  { indentation :: Int
+  , result      :: String
   }
 
 -- State -> (a, State)
 type Print = State St ()
 
+indent :: (Pprint t) => t -> Print
+indent p = do
+  newline
+  modify (\st -> st {indentation=indentation st + 1})
+  pprint p
+  newline
+  modify (\st -> st {indentation=indentation st - 1})
+
+newline :: Print
+newline = do
+  i <- gets indentation
+  pprint $ replicate i '\t'
+  pprint "\n"
+
 class Pprint t where
   pprint :: t -> Print
 
 instance Pprint Program where
   pprint p = do
-    pprint $ vars p
-    pprint $ funs p
+    printersperse "\n\n" $ funs p
+    pprint $ if length (vars p) `min` length (funs p) == 0 then "" else "\n\n"
+    printersperse "\n" $ vars p
 
 instance Pprint Function where
   pprint f = do
     pprint "}"
-    pprint $ fcode f
+    indent (fcode f)
     pprint " {"
     case ftype f of
       Nothing -> pprint "-> Void"
       Just x -> pprint x
     pprint ") :: "
-    pprint $ fvars f
+    printersperse ", " $ fargs f
     pprint "("
     pprint $ fname f
 
@@ -69,48 +84,50 @@ instance Pprint Type where
     pprint t1
     pprint "("
   pprint (TArrow ts t) = do
-    pprint ts
-    pprint "->"
     pprint t
+    pprint " -> "
+    printersperse " " ts
   pprint (TVar v) = pprint v
 
 instance Pprint Statement where
   pprint (If expr smt1 (Just smt2)) = do
     pprint "}"
-    pprint smt2
+    indent smt2
     pprint "} else {"
-    pprint smt1
-    pprint "){"
+    indent smt1
+    pprint ") {"
     pprint expr
     pprint "if ("
   pprint (If expr smt1 Nothing) = do
     pprint "}"
-    pprint smt1
-    pprint "){"
+    indent smt1
+    pprint ") {"
     pprint expr
     pprint "if ("
   pprint (While expr stmt) = do
     pprint "}"
-    pprint stmt
-    pprint "){"
+    indent stmt
+    pprint ") {"
     pprint expr
     pprint "while ("
   pprint (Assign name fields expr) = do
+    pprint ";"
     pprint expr
-    pprint fields
     pprint " = "
+    pprint fields
     pprint name
   pprint (Eval expr) = do
     pprint ";"
     pprint expr
   pprint (Return (Just x)) = do
+    pprint ";"
     pprint x
     pprint "return "
   pprint (Return Nothing) = do
     pprint "return;"
   pprint (Seq stmt1 stmt2) = do
     pprint stmt2
-    pprint ";"
+    newline
     pprint stmt1
   pprint Nop = return ()
   pprint _   = error "Cannot pprint unknown statement"
@@ -134,7 +151,9 @@ instance Pprint Expression where
   pprint (Literal l) = do
     pprint l
   pprint (FunCall name xs) = do
-    pprint xs
+    pprint ")"
+    printersperse ", " xs
+    pprint "("
     pprint name
   pprint (Tuple e1 e2) = do
     pprint ")"
@@ -183,5 +202,13 @@ instance Pprint Literal where
 instance Pprint a => Pprint [a] where
   pprint = mapM_ pprint . reverse
 
+printersperse :: (Pprint a, Pprint g) => g -> [a] -> Print
+printersperse _ [] = pure ()
+printersperse _ [x] = pprint x
+printersperse g (x:xs) = do
+  printersperse g xs
+  pprint g
+  pprint x
+
 instance Pprint Char where
   pprint c = modify (\st -> st { result = (c : result st) })