parser/pretty printer are getting good

app/Expr.hs

@@ -13,6 +13,16 @@ data Expr where
     Pi :: String -> Expr -> Expr -> Expr
     deriving (Show, Eq)
 
+infixl 4 <.>
+
+(<.>) :: Expr -> Expr -> Expr
+(<.>) = App
+
+infixr 2 .->
+
+(.->) :: Expr -> Expr -> Expr
+a .-> b = Pi "" a (incIndices b)
+
 occursFree :: Integer -> Expr -> Bool
 occursFree n (Var k _) = n == k
 occursFree _ Star = False
@@ -26,22 +36,51 @@ occursFree n (Pi _ a b) = occursFree n a || occursFree (n + 1) b
 parenthesize :: String -> String
 parenthesize s = "(" ++ s ++ ")"
 
+collectLambdas :: Expr -> ([(String, Expr)], Expr)
+collectLambdas (Abs x ty body) = ((x, ty) : params, final)
+  where
+    (params, final) = collectLambdas body
+collectLambdas e = ([], e)
+
+collectPis :: Expr -> ([(String, Expr)], Expr)
+collectPis p@(Pi "" _ _) = ([], p)
+collectPis (Pi x ty body) = ((x, ty) : params, final)
+  where
+    (params, final) = collectPis body
+collectPis e = ([], e)
+
+groupParams :: [(String, Expr)] -> [([String], Expr)]
+groupParams = foldr addParam []
+  where
+    addParam :: (String, Expr) -> [([String], Expr)] -> [([String], Expr)]
+    addParam (x, t) [] = [([x], t)]
+    addParam (x, t) l@((xs, s) : rest)
+        | t == s = (x : xs, t) : rest
+        | otherwise = ([x], t) : l
+
+showParamGroup :: ([String], Expr) -> String
+showParamGroup (ids, ty) = parenthesize $ unwords ids ++ " : " ++ pretty ty
+
 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
+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
+helper k (Pi "" t1 t2) = if k > 2 then parenthesize res else res
     where
-    res =
+        res = helper 3 t1 ++ " -> " ++ helper 2 t2
-        if occursFree 0 b
+helper k e@(Pi{}) = if k > 2 then parenthesize res else res
-            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
+    (params, body) = collectPis e
+    grouped = showParamGroup <$> groupParams params
+    res = "∏ " ++ unwords grouped ++ " . " ++ pretty body
+helper k e@(Abs{}) = if k >= 1 then parenthesize res else res
+  where
+    (params, body) = collectLambdas e
+    grouped = showParamGroup <$> groupParams params
+    res = "λ " ++ unwords grouped ++ " . " ++ pretty body
 
 pretty :: Expr -> String
 pretty = helper 0

app/Parser.hs

@@ -1,16 +1,18 @@
+{-# LANGUAGE TupleSections #-}
+
 module Parser where
 
 import Control.Monad
 import Control.Monad.State.Strict
+import Data.Bifunctor (first)
 import Data.Functor.Identity
 import Data.List (elemIndex)
 import Data.List.NonEmpty (NonEmpty ((:|)))
 import qualified Data.List.NonEmpty as NE
-import Expr (Expr (..))
+import Expr (Expr (..), (.->))
 import Text.Megaparsec hiding (State)
 import Text.Megaparsec.Char
 import qualified Text.Megaparsec.Char.Lexer as L
-import Data.Bifunctor (first)
 
 type InnerState = [String]
 
@@ -48,31 +50,44 @@ pVar = label "variable" $ lexeme $ do
 
 defChoice :: NE.NonEmpty String -> Parser ()
 defChoice options = lexeme $ label labelText $ void $ choice $ NE.map string options
-  where labelText = NE.head options
+  where
+    labelText = NE.head options
+
+pParamGroup :: Parser [(String, Expr)]
+pParamGroup = lexeme $ label "parameter group" $ between (char '(') (char ')') $ do
+    idents <- some pIdentifier
+    _ <- defChoice $ ":" :| []
+    ty <- pExpr
+    modify (idents ++)
+    pure $ (,ty) <$> idents
+
+pParams :: Parser [(String, Expr)]
+pParams = concat <$> some pParamGroup
 
 pLAbs :: Parser Expr
 pLAbs = lexeme $ label "λ-abstraction" $ do
-  _ <- defChoice $ "λ" :| ["lambda"]
+    _ <- defChoice $ "λ" :| ["lambda", "fun"]
-  ident <- pIdentifier
+    params <- pParams
-  _ <- defChoice $ ":" :| []
+    _ <- defChoice $ "." :| ["=>", "⇒"]
-  ty <- pExpr
-  _ <- defChoice $ "." :| []
-  modify (ident :)
     body <- pExpr
-  modify $ drop 1
+    modify (drop $ length params)
-  pure $ Abs ident ty body
+    pure $ foldr (uncurry Abs) body params
 
 pPAbs :: Parser Expr
 pPAbs = lexeme $ label "Π-abstraction" $ do
-  _ <- defChoice $ "∏" :| ["Pi"]
+    _ <- defChoice $ "∏" :| ["Pi", "forall", "∀"]
-  ident <- pIdentifier
+    params <- pParams
-  _ <- defChoice $ ":" :| []
+    _ <- defChoice $ "." :| [","]
-  ty <- pExpr
-  _ <- defChoice $ "." :| []
-  modify (ident :)
     body <- pExpr
-  modify $ drop 1
+    modify (drop $ length params)
-  pure $ Pi ident ty body
+    pure $ foldr (uncurry Pi) body params
+
+pArrow :: Parser Expr
+pArrow = lexeme $ label "->" $ do
+    a <- pAppTerm
+    _ <- defChoice $ "->" :| ["→"]
+    b <- pExpr
+    pure $ a .-> b
 
 pApp :: Parser Expr
 pApp = foldl1 App <$> some pTerm
@@ -88,14 +103,17 @@ pTerm =
     lexeme $
         label "term" $
             choice
-        [ between (char '(') (char ')') pExpr,
+                [ between (char '(') (char ')') pExpr
-          pVar,
+                , pVar
-          pStar,
+                , pStar
-          pSquare
+                , pSquare
                 ]
 
+pAppTerm :: Parser Expr
+pAppTerm = lexeme $ pLAbs <|> pPAbs <|> pApp
+
 pExpr :: Parser Expr
-pExpr = lexeme $ pLAbs <|> pPAbs <|> pApp
+pExpr = lexeme $ try pArrow <|> pAppTerm
 
 pAll :: String -> Either String Expr
 pAll input = first errorBundlePretty $ fst $ runIdentity $ runStateT (runParserT pExpr "" input) []