Slightly better pretty printing

2 files changed, 33 insertions, 29 deletions

diff --git a/src/Haskell/AST.hs b/src/Haskell/AST.hs
index e1fd1e4..1e181e2 100644
--- a/src/Haskell/AST.hs
+++ b/src/Haskell/AST.hs
@@ -1,6 +1,7 @@
 module Haskell.AST where
 
 import Data.List
+import Pretty
 
 
 type Name = String
@@ -54,11 +55,8 @@ data Inst = Inst Name Type [Def]
   deriving (Show, Eq)
 
 
-class Pretty a where
-    pretty :: a -> String
-
 instance Pretty AST where
-    pretty (AST tops) = intercalate "\n" (map pretty tops)
+    pretty (AST tops) = Node "" (map pretty tops)
 
 instance Pretty Toplevel where
     pretty (TopDef x) = pretty x
@@ -68,40 +66,46 @@ instance Pretty Toplevel where
     pretty (TopInst x) = pretty x
 
 instance Pretty Def where
-    pretty (Def n e) = n ++ " = " ++ pretty e
+    pretty (Def n e) = Node (n ++ " =") [pretty e]
 
 instance Pretty Expr where
-    pretty (App e as) = "(" ++ intercalate " " (map pretty (e:as)) ++ ")"
-    pretty (Ref n) = n
-    pretty (Num n) = show n
-    pretty (Tup es) = "(" ++ intercalate ", " (map pretty es) ++ ")"
-    pretty (Lam as e) = "(\\" ++ intercalate " " as ++ " -> " ++ pretty e ++ ")"
-    pretty (Case e arms) = "(case " ++ pretty e ++ " of { " ++ intercalate ";" (map go arms) ++ " })"
-      where go (p, e') = pretty p ++ " -> " ++ pretty e'
+    pretty (App e as) = Bracket "(" ")" "" (map pretty (e:as))
+    pretty (Ref n) = Leaf n
+    pretty (Num n) = Leaf (show n)
+    pretty (Tup es) = Bracket "(" ")" "," (map pretty es)
+    pretty (Lam as e) = Bracket "(" ")" "" [Node ("\\" ++ intercalate " " as ++ " ->") [pretty e]]
+    pretty (Case e arms) = Bracket "(" ")" "" [Node ("case " ++ pprintOneline e ++ " of") [Bracket "{" "}" ";" (map go arms)]]
+      where go (p, e') = Node (pprintOneline p ++ " ->") [pretty e']
 
 instance Pretty Pat where
-    pretty PatAny = "_"
-    pretty (PatVar n) = n
-    pretty (PatCon n ps) = "(" ++ n ++ " " ++ intercalate " " (map pretty ps) ++ ")"
-    pretty (PatTup ps) = "(" ++ intercalate ", " (map pretty ps) ++ ")"
+    pretty PatAny = Leaf "_"
+    pretty (PatVar n) = Leaf n
+    pretty (PatCon n ps) = Bracket "(" ")" "" (Leaf n : map pretty ps)
+    pretty (PatTup ps) = Bracket "(" ")" "," (map pretty ps)
 
 instance Pretty Decl where
-    pretty (Decl n t) = n ++ " :: " ++ pretty t
+    pretty (Decl n t) = Node (n ++ " :: ") [pretty t]
 
 instance Pretty Type where
-    pretty (TyTup ts) = "(" ++ intercalate ", " (map pretty ts) ++ ")"
-    pretty TyInt = "Int"
-    pretty (TyFun t u) = "(" ++ pretty t ++ " -> " ++ pretty u ++ ")"
-    pretty (TyRef n as) = "(" ++ n ++ " " ++ intercalate " " (map pretty as) ++ ")"
-    pretty (TyVar n) = n
-    pretty TyVoid = "#Void"
+    pretty (TyTup ts) = Bracket "(" ")" "," (map pretty ts)
+    pretty TyInt = Leaf "Int"
+    pretty (TyFun t u) = Leaf $ pprintOneline t ++ " -> " ++ pprintOneline u
+    pretty (TyRef n as) = Bracket "(" ")" "" (Leaf n : map pretty as)
+    pretty (TyVar n) = Leaf n
+    pretty TyVoid = Leaf "#Void"
 
 instance Pretty Data where
-    pretty (Data n as cs) = "data " ++ n ++ " " ++ intercalate " " as ++ " = " ++ intercalate " | " (map go cs)
-      where go (m, ts) = m ++ " " ++ intercalate " " (map pretty ts)
+    pretty (Data n as cs) = Node ("data " ++ n ++ " " ++ intercalate " " as ++ " =") [Bracket "" "" "|" (map go cs)]
+      where go (m, ts) = Node m (map pretty ts)
 
 instance Pretty Class where
-    pretty (Class n as ds) = "class " ++ n ++ " " ++ intercalate " " as ++ " where { " ++ intercalate " ; " (map pretty ds) ++ "}"
+    pretty (Class n as ds) = Node ("class " ++ n ++ " " ++ intercalate " " as ++ " where") [Bracket "{" "}" ";" (map pretty ds)]
 
 instance Pretty Inst where
-    pretty (Inst n t ds) = "instance " ++ n ++ " " ++ pretty t ++ " where { " ++ intercalate " ; " (map pretty ds) ++ " }"
+    pretty (Inst n t ds) = Node ("instance " ++ n ++ " " ++ pprintOneline t ++ " where") [Bracket "{" "}" ";" (map pretty ds)]
+
+
+mapInit :: (a -> a) -> [a] -> [a]
+mapInit _ [] = []
+mapInit _ [x] = [x]
+mapInit f (x:y:zs) = f x : mapInit f (y:zs)
diff --git a/src/Haskell/Env.hs b/src/Haskell/Env.hs
index 6b74221..aca2367 100644
--- a/src/Haskell/Env.hs
+++ b/src/Haskell/Env.hs
@@ -1,9 +1,9 @@
 module Haskell.Env where
 
 import Control.Monad
-import Data.List
 import qualified Data.Map.Strict as Map
 import Haskell.AST
+import Pretty
 
 
 data Env = Env { eDefs :: Map.Map Name Expr }
@@ -11,7 +11,7 @@ data Env = Env { eDefs :: Map.Map Name Expr }
 
 instance Pretty Env where
     pretty (Env defs) =
-        intercalate "\n" [n ++ " = " ++ pretty e | (n, e) <- Map.assocs defs]
+        Node "" [Node (n ++ " =") [pretty e] | (n, e) <- Map.assocs defs]
 
 emptyEnv :: Env
 emptyEnv = Env Map.empty