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 = helper 3 t1 ++ " -> " ++ helper 2 t2
+helper k e@(Pi{}) = 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
+    (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
-    res = "λ" ++ x ++ " : " ++ helper 0 ty ++ " . " ++ helper 0 b
+    (params, body) = collectLambdas e
+    grouped = showParamGroup <$> groupParams params
+    res = "λ " ++ unwords grouped ++ " . " ++ pretty body
 
 pretty :: Expr -> String
 pretty = helper 0

app/Parser.hs

@@ -1,78 +1,93 @@
+{-# 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 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]
 
 data CustomErrors = UnboundVariable String [String] deriving (Eq, Ord, Show)
 
 instance ShowErrorComponent CustomErrors where
-  showErrorComponent (UnboundVariable var bound) =
-    "Unbound variable: " ++ var ++ ". Did you mean one of: " ++ unwords bound ++ "?"
+    showErrorComponent (UnboundVariable var bound) =
+        "Unbound variable: " ++ var ++ ". Did you mean one of: " ++ unwords bound ++ "?"
 
 type Parser = ParsecT CustomErrors String (State InnerState)
 
 skipSpace :: Parser ()
 skipSpace =
-  L.space
-    space1
-    (L.skipLineComment "--")
-    (L.skipBlockCommentNested "(*" "*)")
+    L.space
+        space1
+        (L.skipLineComment "--")
+        (L.skipBlockCommentNested "(*" "*)")
 
 lexeme :: Parser a -> Parser a
 lexeme = L.lexeme skipSpace
 
 pIdentifier :: Parser String
 pIdentifier = label "identifier" $ lexeme $ do
-  firstChar <- letterChar <|> char '_'
-  rest <- many $ alphaNumChar <|> char '_'
-  return $ firstChar : rest
+    firstChar <- letterChar <|> char '_'
+    rest <- many $ alphaNumChar <|> char '_'
+    return $ firstChar : rest
 
 pVar :: Parser Expr
 pVar = label "variable" $ lexeme $ do
-  var <- pIdentifier
-  binders <- get
-  case elemIndex var binders of
-    Just i -> return $ Var (fromIntegral i) var
-    Nothing -> customFailure $ UnboundVariable var binders
+    var <- pIdentifier
+    binders <- get
+    case elemIndex var binders of
+        Just i -> return $ Var (fromIntegral i) var
+        Nothing -> customFailure $ UnboundVariable var binders
 
 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"]
-  ident <- pIdentifier
-  _ <- defChoice $ ":" :| []
-  ty <- pExpr
-  _ <- defChoice $ "." :| []
-  modify (ident :)
-  body <- pExpr
-  modify $ drop 1
-  pure $ Abs ident ty body
+    _ <- defChoice $ "λ" :| ["lambda", "fun"]
+    params <- pParams
+    _ <- defChoice $ "." :| ["=>", "⇒"]
+    body <- pExpr
+    modify (drop $ length params)
+    pure $ foldr (uncurry Abs) body params
 
 pPAbs :: Parser Expr
 pPAbs = lexeme $ label "Π-abstraction" $ do
-  _ <- defChoice $ "∏" :| ["Pi"]
-  ident <- pIdentifier
-  _ <- defChoice $ ":" :| []
-  ty <- pExpr
-  _ <- defChoice $ "." :| []
-  modify (ident :)
-  body <- pExpr
-  modify $ drop 1
-  pure $ Pi ident ty body
+    _ <- defChoice $ "∏" :| ["Pi", "forall", "∀"]
+    params <- pParams
+    _ <- defChoice $ "." :| [","]
+    body <- pExpr
+    modify (drop $ length params)
+    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
@@ -85,17 +100,20 @@ pSquare = Square <$ defChoice ("□" :| ["[]"])
 
 pTerm :: Parser Expr
 pTerm =
-  lexeme $
-    label "term" $
-      choice
-        [ between (char '(') (char ')') pExpr,
-          pVar,
-          pStar,
-          pSquare
-        ]
+    lexeme $
+        label "term" $
+            choice
+                [ between (char '(') (char ')') pExpr
+                , pVar
+                , pStar
+                , 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) []