From e9e388ba050175a65c61cb641e6ea813eb0bee6d Mon Sep 17 00:00:00 2001
From: William Ball <williampi103@gmail.com>
Date: Mon, 11 Nov 2024 14:34:55 -0800
Subject: [PATCH] greatly improved pretty printer

---
 app/Expr.hs | 32 +++++++++++++++++++++-----------
 app/Main.hs |  9 +++++----
 2 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/app/Expr.hs b/app/Expr.hs
index 035b294..0708b66 100644
--- a/app/Expr.hs
+++ b/app/Expr.hs
@@ -23,18 +23,28 @@ occursFree n (Pi _ a b) = occursFree n a || occursFree (n + 1) b
 
 {- --------------------- PRETTY PRINTING ----------------------------- -}
 
--- TODO : store parsed identifiers for better printing
+parenthesize :: String -> String
+parenthesize s = "(" ++ s ++ ")"
+
+helper :: Integer -> Expr -> String
+helper _ (Var _ s) = s
+helper _ Star = "*"
+helper _ Square = "□"
+helper k (App e1 e2) = if k >= 3 then parenthesize res else res
+  where
+    res = helper 3 e1 ++ " " ++ helper 4 e2
+helper k (Pi x ty b) = if k >= 2 then parenthesize res else res
+  where
+    res =
+        if occursFree 0 b
+            then "∏" ++ x ++ " : " ++ helper 0 ty ++ " . " ++ helper 0 b
+            else helper 3 ty ++ " -> " ++ helper 2 b
+helper k (Abs x ty b) = if k >= 1 then parenthesize res else res
+  where
+    res = "λ" ++ x ++ " : " ++ helper 0 ty ++ " . " ++ helper 0 b
+
 pretty :: Expr -> String
-pretty (Var _ s) = s
-pretty Star = "*"
-pretty Square = "□"
-pretty (App e1 e2) = "(" ++ pretty e1 ++ " " ++ pretty e2 ++ ")"
-pretty (Abs x ty b) = "(λ" ++ x ++ " : " ++ pretty ty ++ " . " ++ pretty b ++ ")"
-pretty (Pi x ty b) = 
-    if occursFree 0 b then
-        "(∏" ++ x ++ " : " ++ pretty ty ++ " . " ++ pretty b ++ ")"
-    else
-        "(" ++ pretty ty ++ " -> " ++ pretty b ++ ")"
+pretty = helper 0
 
 {- --------------- ACTUAL MATH STUFF ---------------- -}
 
diff --git a/app/Main.hs b/app/Main.hs
index 06381d4..1cabb2c 100644
--- a/app/Main.hs
+++ b/app/Main.hs
@@ -3,6 +3,7 @@ module Main where
 import Expr
 import Parser
 import System.IO
+
 -- import Check
 
 main :: IO ()
@@ -12,8 +13,8 @@ main = do
     input <- getLine
     case pAll input of
         Left err -> putStrLn err
-        Right expr -> putStrLn $ pretty expr
-        -- Right expr -> case findType [] expr of
-        --     Just ty -> putStrLn $ pretty expr ++ " : " ++ pretty ty
-        --     Nothing -> putStrLn $ "Unable to find type for " ++ pretty expr ++ "!"
+        Right expr -> putStrLn (pretty expr)
+    -- Right expr -> case findType [] expr of
+    --     Just ty -> putStrLn $ pretty expr ++ " : " ++ pretty ty
+    --     Nothing -> putStrLn $ "Unable to find type for " ++ pretty expr ++ "!"
     main