1 files changed, 94 insertions, 0 deletions
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