From fba46296db5aabb475df321f166770b0c311f7da Mon Sep 17 00:00:00 2001
From: crumbtoo <crumb@disroot.org>
Date: Thu, 8 Feb 2024 11:40:13 -0700
Subject: [PATCH] ppr typesigs

---
 rlp.cabal          |  1 +
 src/Core/Syntax.hs | 36 +++++++++++++++++++++++++++++++++---
 src/Data/Pretty.hs |  9 +++++++++
 3 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/rlp.cabal b/rlp.cabal
index 1f18e4d..7ed9477 100644
--- a/rlp.cabal
+++ b/rlp.cabal
@@ -78,6 +78,7 @@ library
                     , microlens-pro ^>=0.2.0
                     , effectful-core ^>=2.3.0.0
                     , deriving-compat ^>=0.6.0
+                    , these >=0.2 && <2.0
 
     hs-source-dirs:   src
     default-language: GHC2021
diff --git a/src/Core/Syntax.hs b/src/Core/Syntax.hs
index f95163e..ebc6dca 100644
--- a/src/Core/Syntax.hs
+++ b/src/Core/Syntax.hs
@@ -57,6 +57,8 @@ import Data.HashMap.Strict          qualified as H
 import Data.Hashable
 import Data.Text                    qualified as T
 import Data.Char
+import Data.These
+import Data.Bifoldable              (bifoldr)
 import GHC.Generics                 (Generic, Generically(..))
 -- Lift instances for the Core quasiquoters
 import Language.Haskell.TH.Syntax   (Lift)
@@ -220,11 +222,39 @@ instance HasLHS (Binding b) (Binding b) b b where
 
 -- TODO: print type sigs with corresponding scdefs
 -- TODO: emit pragmas for datatags
-instance (Pretty b) => Pretty (Program b) where
-    pretty = vsepOf (programScDefs . each . to pretty)
+instance (Hashable b, Pretty b) => Pretty (Program b) where
+    -- pretty = vsepOf (programScDefs . each . to pretty)
+    pretty = vlinesOf (programJoinedDefs . to prettyGroup) where
+        programJoinedDefs :: Fold (Program b) (These (b, Type) (ScDef b))
+        programJoinedDefs = folding $ \p ->
+            foldMapOf programTypeSigs thing1 p
+            `u` foldMapOf programScDefs thing2 p
+          where u = H.unionWith unionThese
+
+        thing1 = ifoldMap @b @(HashMap b)
+            (\n t -> H.singleton n (This (n,t)))
+        thing2 = foldMap $ \sc ->
+            H.singleton (sc ^. _lhs . _1) (That sc)
+
+        prettyGroup :: These (b, Type) (ScDef b) -> Doc
+        prettyGroup = bifoldr ($$) ($$) mempty . bimap prettyTySig pretty
+
+        prettyTySig (n,t) = hsep [ttext n, "::", pretty t]
+
+        unionThese :: forall a b. These a b -> These a b -> These a b
+        unionThese (This a) (That b) = These a b
+        unionThese (That b) (This a) = These a b
+        unionThese (These a b) _     = These a b
+
+instance Pretty Type where
+    prettyPrec _ (TyVar n) = ttext n
+    prettyPrec _ TyFun = "(->)"
+    prettyPrec _ (TyCon n) = ttext n
+    prettyPrec p (TyApp f x) = maybeParens (p>0) $
+        prettyPrec 0 f <+> prettyPrec 1 x
 
 instance (Pretty b) => Pretty (ScDef b) where
-    pretty sc = hsep [name, as, "=", hang empty 1 e]
+    pretty sc = hsep [name, as, "=", hang empty 1 e, ";"]
         where
             name = ttext $ sc ^. _lhs . _1
             as = sc & hsepOf (_lhs . _2 . each . to ttext)
diff --git a/src/Data/Pretty.hs b/src/Data/Pretty.hs
index f5b1b4d..f16c319 100644
--- a/src/Data/Pretty.hs
+++ b/src/Data/Pretty.hs
@@ -3,6 +3,8 @@ module Data.Pretty
     , ttext
     -- * Pretty-printing lens combinators
     , hsepOf, vsepOf
+    , vcatOf
+    , vlinesOf
     , module Text.PrettyPrint
     , maybeParens
     )
@@ -49,3 +51,10 @@ hsepOf l = foldrOf l (<+>) mempty
 vsepOf :: Getting (Endo Doc) s Doc -> s -> Doc
 vsepOf l = foldrOf l ($+$) mempty
 
+vcatOf :: Getting (Endo Doc) s Doc -> s -> Doc
+vcatOf l = foldrOf l ($$) mempty
+
+vlinesOf :: Getting (Endo Doc) s Doc -> s -> Doc
+vlinesOf l = foldrOf l (\a b -> a $+$ "" $+$ b) mempty
+-- hack(?) to separate chunks with a blankline
+