more parser goodness

lib/Check.hs

@@ -33,7 +33,6 @@ findType g (Var x n) = do
     validateType g t
     pure t
 findType _ (Free n) = envLookupTy n
-findType _ (Axiom n) = envLookupTy n
 findType g e@(App m n) = do
     (a, b) <- findType g m >>= matchPi m
     a' <- findType g n

lib/Expr.hs

@@ -3,7 +3,6 @@ module Expr where
 data Expr where
     Var :: Text -> Integer -> Expr
     Free :: Text -> Expr
-    Axiom :: Text -> Expr
     Level :: Integer -> Expr
     App :: Expr -> Expr -> Expr
     Abs :: Text -> Expr -> Expr -> Expr
@@ -14,7 +13,6 @@ data Expr where
 instance Eq Expr where
     (Var _ n) == (Var _ m) = n == m
     (Free s) == (Free t) = s == t
-    (Axiom a) == (Axiom b) = a == b
     (Level i) == (Level j) = i == j
     (App e1 e2) == (App f1 f2) = e1 == f1 && e2 == f2
     (Abs _ t1 b1) == (Abs _ t2 b2) = t1 == t2 && b1 == b2
@@ -25,7 +23,6 @@ instance Eq Expr where
 occursFree :: Integer -> Expr -> Bool
 occursFree n (Var _ k) = n == k
 occursFree _ (Free _) = False
-occursFree _ (Axiom _) = False
 occursFree _ (Level _) = False
 occursFree n (App a b) = on (||) (occursFree n) a b
 occursFree n (Abs _ a b) = occursFree n a || occursFree (n + 1) b
@@ -37,7 +34,6 @@ shiftIndices d c (Var x k)
     | k >= c = Var x (k + d)
     | otherwise = Var x k
 shiftIndices _ _ (Free s) = Free s
-shiftIndices _ _ (Axiom s) = Axiom s
 shiftIndices _ _ (Level i) = Level i
 shiftIndices d c (App m n) = App (shiftIndices d c m) (shiftIndices d c n)
 shiftIndices d c (Abs x m n) = Abs x (shiftIndices d c m) (shiftIndices d (c + 1) n)
@@ -109,7 +105,6 @@ showBinding (ident, params, val) =
 helper :: Integer -> Expr -> Text
 helper _ (Var s _) = s
 helper _ (Free s) = s
-helper _ (Axiom s) = s
 helper _ (Level i)
     | i == 0 = "*"
     | otherwise = "*" <> show i

lib/IR.hs

@@ -8,7 +8,6 @@ type Param = (Text, IRExpr)
 
 data IRExpr
     = Var {_varName :: Text}
-    | Axiom
     | Level {_level :: Integer}
     | App
         { _appFunc :: IRExpr
@@ -34,11 +33,17 @@ data IRExpr
 
 makeLenses ''IRExpr
 
-data IRDef = Def
+data IRDef
+    = Def
         { _defName :: Text
         , _defParams :: [Param]
         , _defAscription :: Maybe IRExpr
         , _defBody :: IRExpr
         }
+    | Axiom
+        { _axiomName :: Text
+        , _axiomParams :: [Param]
+        , _axiomAscription :: IRExpr
+        }
 
 type IRProgram = [IRDef]

lib/Parser.hs

@@ -3,12 +3,13 @@
 module Parser where
 
 import Control.Monad.Except
+import Data.Char
 import Data.List (foldl1)
 import qualified Data.Text as T
 import Errors (Error (..))
 import IR
 import Preprocessor
-import Text.Megaparsec (Parsec, ShowErrorComponent (..), between, choice, chunk, errorBundlePretty, label, runParser, try)
+import Text.Megaparsec (Parsec, ShowErrorComponent (..), between, choice, chunk, errorBundlePretty, label, runParser, satisfy, try)
 import Text.Megaparsec.Char
 import qualified Text.Megaparsec.Char.Lexer as L
 
@@ -34,17 +35,19 @@ eat :: Text -> Parser ()
 eat = void . lexeme . chunk
 
 keywords :: [Text]
-keywords = ["forall", "let", "in", "end", "fun"]
+keywords = ["forall", "let", "in", "end", "fun", "def", "axiom"]
+
+reservedChars :: [Char]
+reservedChars = "();:"
 
 pIdentifier :: Parser Text
 pIdentifier = try $ label "identifier" $ lexeme $ do
-    firstChar <- letterChar <|> char '_'
-    rest <- many $ alphaNumChar <|> char '_'
-    let ident = T.pack (firstChar : rest)
-    when (ident `elem` keywords) $
-        fail $
-            "Reserved word: " ++ T.unpack ident
+    chars <- many $ satisfy isAllowed
+    let ident = T.pack chars
+    when (ident `elem` keywords) $ fail $ "Reserved word: " ++ T.unpack ident
     pure ident
+  where
+    isAllowed c = isLetter c || isNumber c || c == '_' || isSymbol c && c `notElem` reservedChars
 
 pVar :: Parser IRExpr
 pVar = label "variable" $ lexeme $ Var <$> pIdentifier
@@ -86,7 +89,7 @@ pBinding = lexeme $ label "binding" $ do
     eat "("
     ident <- pIdentifier
     params <- pManyParams
-    ascription <- pAscription
+    ascription <- optional pAscription
     eat ":="
     value <- pIRExpr
     eat ")"
@@ -126,20 +129,31 @@ pNumSort = lexeme $ label "sort" $ eat "*" >> Level <$> L.decimal
 pSort :: Parser IRExpr
 pSort = try pNumSort <|> pStar
 
-pAxiom :: Parser IRExpr
-pAxiom = Axiom <$ eat "axiom"
-
-pIRDef :: Parser IRDef
-pIRDef = lexeme $ label "definition" $ do
+pAxiom :: Parser IRDef
+pAxiom = lexeme $ label "axiom" $ do
     skipSpace
+    eat "def"
     ident <- pIdentifier
     params <- pManyParams
-    ascription <- fmap (flip (foldr (uncurry Pi)) params) <$> pAscription
+    ascription <- fmap (flip (foldr (uncurry Pi)) params) pAscription
+    eat ";"
+    pure $ Axiom ident params ascription
+
+pDef :: Parser IRDef
+pDef = lexeme $ label "definition" $ do
+    skipSpace
+    eat "def"
+    ident <- pIdentifier
+    params <- pManyParams
+    ascription <- fmap (flip (foldr (uncurry Pi)) params) <$> optional pAscription
     eat ":="
     body <- pIRExpr
     eat ";"
     pure $ Def ident params ascription body
 
+pIRDef :: Parser IRDef
+pIRDef = pDef <|> pAxiom
+
 pTerm :: Parser IRExpr
 pTerm =
     lexeme $
@@ -147,7 +161,6 @@ pTerm =
             choice
                 [ between (char '(') (char ')') pIRExpr
                 , pSort
-                , pAxiom
                 , pVar
                 ]
 
@@ -164,8 +177,8 @@ pAppTerm =
 pIRExpr :: Parser IRExpr
 pIRExpr = lexeme $ try pArrow <|> pAppTerm
 
-pAscription :: Parser (Maybe IRExpr)
-pAscription = lexeme $ optional $ try $ eat ":" >> label "type" pIRExpr
+pAscription :: Parser IRExpr
+pAscription = lexeme $ try $ eat ":" >> label "type" pIRExpr
 
 pIRProgram :: Parser IRProgram
 pIRProgram = many pIRDef