Slightly better pretty printing

4 files changed, 141 insertions, 42 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
diff --git a/src/Main.hs b/src/Main.hs
index 61dc3bc..faf21f7 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -6,6 +6,7 @@ import Haskell.AST
 import Haskell.Env
 import Haskell.Rewrite
 import Haskell.Parser
+import Pretty
 import System.Environment
 import System.Exit
 import System.IO
@@ -95,17 +96,17 @@ data UserCmd = UCCmd Cmd
   deriving (Show)
 
 instance Pretty Action where
-    pretty (Action cmd target) = "action " ++ target ++ " " ++ pretty cmd
+    pretty (Action cmd target) = Node ("action " ++ target) [pretty cmd]
 
 instance Pretty Cmd where
-    pretty (CRewrite name) = "rew " ++ name
-    pretty CBeta = "beta"
-    pretty CEta = "eta"
-    pretty CCase = "case"
-    pretty CCaseForce = "case!"
-    pretty CEtaCase = "etacase"
-    pretty CCaseCase = "casecase"
-    pretty (CRepeat cmd) = "repeat " ++ pretty cmd
+    pretty (CRewrite name) = Leaf ("rew " ++ name)
+    pretty CBeta = Leaf "beta"
+    pretty CEta = Leaf "eta"
+    pretty CCase = Leaf "case"
+    pretty CCaseForce = Leaf "case!"
+    pretty CEtaCase = Leaf "etacase"
+    pretty CCaseCase = Leaf "casecase"
+    pretty (CRepeat cmd) = Node "repeat" [pretty cmd]
 
 topEnv :: Context -> Env
 topEnv (Context [] env) = env
@@ -134,7 +135,7 @@ apply ctx act@(Action cmd target) = case cmd of
     CRewrite name -> genTransform (\env -> liftM2 (,) (get env name) (get env target))
                                   (\(repl, expr) -> normalise $ rewrite name repl expr)
     CRepeat subcmd -> case exprTransformer subcmd of
-                          Nothing -> Left $ "Cannot repeat '" ++ pretty subcmd ++ "'"
+                          Nothing -> Left $ "Cannot repeat '" ++ pprint 80 subcmd ++ "'"
                           Just trans -> transform (repeatTrans (normalise . trans))
     _ -> transform (normalise . fromJust (exprTransformer cmd))
   where
@@ -165,13 +166,13 @@ applyUserCmd appstate = \case
         in case cStack ctx of
                  [] -> return (Left "Empty environment stack")
                  (_:stk) -> return (Right (Just appstate { asCtx = ctx { cStack = stk } }))
-    CShowEnv -> putStrLn (pretty (topEnv (asCtx appstate))) >> return (Right (Just appstate))
+    CShowEnv -> putStrLn (pprint 80 (topEnv (asCtx appstate))) >> return (Right (Just appstate))
     CUnFocus -> return (Right (Just appstate { asFocus = Nothing }))
     CFocus name -> if envContains (topEnv (asCtx appstate)) name
                        then return (Right (Just appstate { asFocus = Just name }))
                        else return (Left "Name doesn't exist")
     CLog -> do
-        mapM_ (putStrLn . pretty) $ map fst (reverse $ cStack (asCtx appstate))
+        mapM_ (putStrLn . pprint 80) $ map fst (reverse $ cStack (asCtx appstate))
         return (Right (Just appstate))
     CHelp -> putStrLn usageString >> return (Right (Just appstate))
 
@@ -199,7 +200,7 @@ interface appstate = do
         Just n  ->
             case get (topEnv (asCtx appstate')) n of
                 Left _ -> return ()
-                Right expr -> putStrLn $ pretty (Def n expr)
+                Right expr -> putStrLn $ pprint 80 (Def n expr)
 
     putStr "> " >> hFlush stdout
 
diff --git a/src/Pretty.hs b/src/Pretty.hs
new file mode 100644
index 0000000..935bc03
--- /dev/null
+++ b/src/Pretty.hs
@@ -0,0 +1,94 @@
+module Pretty(PrTree(..), Pretty(..), pprint, pprintOneline) where
+
+import Data.Char
+import Data.List
+import Data.Ord
+
+
+data PrTree = Leaf String
+            | Node String [PrTree]
+            -- Bracket "(" ")" "," [1,2,3] => (1, 2, 3)
+            | Bracket String String String [PrTree]
+
+class Pretty a where
+    pretty :: a -> PrTree
+
+data Budget = Budget { bNode :: !Int, bBracket :: !Int }
+
+decNode :: Budget -> Budget
+decNode b = b { bNode = bNode b - 1 }
+
+decBracket :: Budget -> Budget
+decBracket b = b { bBracket = bBracket b - 1 }
+
+printingBudget :: Budget
+printingBudget = Budget 5 5
+
+pprint :: Pretty a => Int -> a -> String
+pprint w val = pprintPr printingBudget w (pretty val)
+
+pprintOneline :: Pretty a => a -> String
+pprintOneline = pprint maxBound
+
+pprintPr :: Budget -> Int -> PrTree -> String
+pprintPr budget w val = intercalate "\n" $ pprintPrX budget w 0 val
+
+pprintPrOneline :: Budget -> PrTree -> String
+pprintPrOneline budget = pprintPr budget maxBound
+
+pprintPrX :: Budget -> Int -> Int -> PrTree -> [String]
+pprintPrX _      _   _ (Leaf str) = [str]
+pprintPrX budget wid x (Node "" subs) = map (pprintPrOneline (decNode budget)) subs
+pprintPrX budget wid x (Node prefix subs) =
+    let subrenders = map (pprintPrX (decNode budget) wid (x + indentWid)) subs
+    in if any ((> 1) . length) subrenders
+           then prefix : map (spaces indentWid ++) (concat subrenders)
+           else [prefix ++ " " ++ intercalate " " (concat subrenders)]
+-- pprintPrX budget wid x (Node prefix subs) =
+--     let prefix' = case prefix of
+--                       "" -> ""
+--                       _  -> prefix ++ " "
+--     in chooseBest (bNode budget) wid
+--         [[prefix' ++ intercalate " " (map (pprintPrOneline (decNode budget)) subs)]
+--         ,let subrender = concatMap (pprintPrX (decNode budget) wid (x + length prefix')) subs
+--          in case subrender of
+--                 [] -> [prefix']
+--                 (ln:lns) -> (prefix' ++ ln) : [spaces (length prefix') ++ l | l <- lns]
+--         ,let subrender = concatMap (pprintPrX (decNode budget) wid (x + indentWid)) subs
+--          in prefix : map (spaces indentWid ++) subrender]
+pprintPrX budget wid x (Bracket beg end sep subs) = chooseBest (bBracket budget) wid
+    [[beg ++ intercalate (sep ++ " ") (map (pprintPrOneline (decBracket budget)) subs) ++ end]
+    ,let subrender = concatMap (fmapLast (++ sep) . pprintPrX (decBracket budget) wid (x + length beg)) subs
+     in case subrender of
+            [] -> [beg ++ end]
+            [ln] -> [beg ++ ln ++ end]
+            (ln:lns) -> (beg ++ ln) : fmapLast (++ end) [spaces (length beg) ++ l | l <- lns]
+    ,let subrender = concatMap (fmapLast (++ sep) . pprintPrX (decBracket budget) wid (x + indentWid)) subs
+     in case subrender of
+            [] -> [beg ++ end]
+            lns -> [trimRight beg] ++ [spaces indentWid ++ l | l <- lns] ++ [trimLeft end]]
+
+chooseBest :: Int -> Int -> [[String]] -> [String]
+chooseBest n _ | n <= 0 = head
+chooseBest _ wid = argmin $ \lns -> (sum [max 0 (length l - wid) | l <- lns]
+                                    ,length lns)
+
+indentWid :: Int
+indentWid = 4
+
+argmin :: (Show a, Show b, Ord b) => (a -> b) -> [a] -> a
+argmin f l = fst $ minimumBy (comparing snd) [(x, f x) | x <- l]
+
+spaces :: Int -> String
+spaces n = replicate n ' '
+
+fmapLast :: (a -> a) -> [a] -> [a]
+fmapLast _ [] = []
+fmapLast f [x] = [f x]
+fmapLast f (x:y:zs) = x : fmapLast f (y:zs)
+
+trimLeft :: String -> String
+trimLeft = dropWhile isSpace
+
+trimRight :: String -> String
+trimRight = reverse . trimLeft . reverse