From d265a423b7e62506781774c34f14632adb94f7d8 Mon Sep 17 00:00:00 2001
From: crumbtoo <crumb@disroot.org>
Date: Thu, 9 Nov 2023 18:59:00 -0700
Subject: [PATCH] pretty printer

---
 src/Core.hs        | 20 +++++++++++++++++---
 src/Data/Pretty.hs | 46 ++++++++++++++++++++--------------------------
 2 files changed, 37 insertions(+), 29 deletions(-)

diff --git a/src/Core.hs b/src/Core.hs
index f821841..79248e5 100644
--- a/src/Core.hs
+++ b/src/Core.hs
@@ -1,8 +1,9 @@
-{-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE PatternSynonyms, OverloadedStrings #-}
 module Core where
 ----------------------------------------------------------------------------------
 import Data.Coerce
 import Data.Pretty
+import Data.List            (intersperse)
 ----------------------------------------------------------------------------------
 
 data Expr = Var Name
@@ -27,7 +28,7 @@ pattern k := v = Binding k v
 
 data Rec = Rec
          | NonRec
-         deriving Show
+         deriving (Show, Eq)
 
 data Alter = Alter Int [Name] Expr
     deriving Show
@@ -41,7 +42,20 @@ newtype Program = Program [ScDef]
 ----------------------------------------------------------------------------------
 
 instance Pretty Expr where
-    prettyPrec _ (Var k) = iStr k
+    prettyPrec _ (Var k)      = IStr k
+    prettyPrec _ (IntP n)     = IStr $ show n
+    prettyPrec _ (Con _ _)    = undefined
+    prettyPrec _ (Let r bs e) =
+        IStr (if r == Rec then "letrec " else "let ")
+        <> binds <> IBreak
+        <> "in " <> prettyPrec 0 e
+        where
+            binds = mconcat (fmap f (init bs))
+                 <> IIndent (prettyPrec 0 $ last bs)
+            f b = IIndent $ prettyPrec 0 b <> IBreak
+
+instance Pretty Binding where
+    prettyPrec _ (k := v) = IStr k <> " = " <> prettyPrec 0 v
 
 ----------------------------------------------------------------------------------
 
diff --git a/src/Data/Pretty.hs b/src/Data/Pretty.hs
index 52c9b2b..c7374dc 100644
--- a/src/Data/Pretty.hs
+++ b/src/Data/Pretty.hs
@@ -1,24 +1,30 @@
 module Data.Pretty
     ( Pretty(..)
-    , ISeq
-    , iNil
-    , iStr
-    , iAppend
+    , ISeq(..)
+    , iBracket
     )
     where
 ----------------------------------------------------------------------------------
+import Data.String      (IsString(..))
+----------------------------------------------------------------------------------
 
 class Pretty a where
     pretty :: a -> String
     prettyPrec :: Int -> a -> ISeq
 
     pretty = squash . prettyPrec 0
-    prettyPrec _ a = iBracket (iStr $ pretty a)
+    prettyPrec _ a = iBracket (IStr $ pretty a)
+    {-# MINIMAL pretty | prettyPrec #-}
 
 data ISeq where
     INil :: ISeq
     IStr :: String -> ISeq
     IAppend :: ISeq -> ISeq -> ISeq
+    IIndent :: ISeq -> ISeq
+    IBreak :: ISeq
+
+instance IsString ISeq where
+    fromString = IStr
 
 instance Semigroup ISeq where
     (<>) = IAppend
@@ -27,28 +33,16 @@ instance Monoid ISeq where
     mempty = INil
 
 squash :: ISeq -> String
-squash = flatten . pure
+squash a = flatten 0 [(a,0)]
 
-flatten :: [ISeq] -> String
-flatten (INil : ss)        = flatten ss
-flatten (IStr s : ss)      = s ++ flatten ss
-flatten (IAppend r s : ss) = flatten (r : s : ss)
-
-iNil :: ISeq
-iNil = INil
-
-iStr :: String -> ISeq
-iStr = IStr
-
-iAppend  :: ISeq -> ISeq -> ISeq
-iAppend = IAppend
-
-iIndent :: ISeq -> ISeq
-iIndent = id
-
-iBreak :: ISeq
-iBreak = iStr "\n"
+flatten :: Int -> [(ISeq, Int)] -> String
+flatten _ []                      = ""
+flatten c ((INil,        i) : ss) = flatten c ss
+flatten c ((IStr s,      i) : ss) = s ++ flatten (c + length s) ss
+flatten c ((IAppend r s, i) : ss) = flatten c ((r,i) : (s,i) : ss)
+flatten _ ((IBreak,      i) : ss) = '\n' : replicate i ' ' ++ flatten i ss
+flatten c ((IIndent s,   i) : ss) = flatten c ((s,c) : ss)
 
 iBracket :: ISeq -> ISeq
-iBracket s = iStr "(" `iAppend` s `iAppend` iStr ")"
+iBracket s = IStr "(" <> s <> IStr ")"