From 9f5c3081317913b26ed0fca993935d7ad8fab53c Mon Sep 17 00:00:00 2001
From: William Ball <williampi103@gmail.com>
Date: Sun, 1 Dec 2024 15:28:57 -0800
Subject: [PATCH] IR success!

---
 lib/Elaborator.hs | 24 +++++++++++++++++++-----
 lib/Parser.hs     | 29 ++++++++++++-----------------
 2 files changed, 31 insertions(+), 22 deletions(-)

diff --git a/lib/Elaborator.hs b/lib/Elaborator.hs
index 2299470..0af26e6 100644
--- a/lib/Elaborator.hs
+++ b/lib/Elaborator.hs
@@ -8,22 +8,36 @@ import qualified IR as I
 
 type Binders = [Text]
 
+saveBinders :: State Binders a -> State Binders a
+saveBinders action = do
+    binders <- get
+    res <- action
+    put binders
+    pure res
+
 elaborate :: IRExpr -> Expr
 elaborate ir = evalState (elaborate' ir) []
-    where
+  where
     elaborate' :: IRExpr -> State Binders Expr
     elaborate' (I.Var n) = do
         binders <- get
         pure $ E.Var n . fromIntegral <$> elemIndex n binders ?: E.Free n
     elaborate' (I.Level level) = pure $ E.Level level
     elaborate' (I.App m n) = E.App <$> elaborate' m <*> elaborate' n
-    elaborate' (I.Abs x t b) = do
+    elaborate' (I.Abs x t b) = saveBinders $ do
         t' <- elaborate' t
         modify (x :)
         E.Abs x t' <$> elaborate' b
-    elaborate' (I.Pi x t b) = do
+    elaborate' (I.Pi x t b) = saveBinders $ do
         t' <- elaborate' t
         modify (x :)
         E.Pi x t' <$> elaborate' b
-    elaborate' (I.Let name Nothing val body) = E.Let name Nothing <$> elaborate' val <*> elaborate' body
-    elaborate' (I.Let name (Just t) val body) = E.Let name . Just <$> elaborate' t <*> elaborate' val <*> elaborate' body
+    elaborate' (I.Let name Nothing val body) = saveBinders $ do
+        val' <- elaborate' val
+        modify (name :)
+        E.Let name Nothing val' <$> elaborate' body
+    elaborate' (I.Let name (Just ty) val body) = saveBinders $ do
+        val' <- elaborate' val
+        ty' <- elaborate' ty
+        modify (name :)
+        E.Let name (Just ty') val' <$> elaborate' body
diff --git a/lib/Parser.hs b/lib/Parser.hs
index 0f31292..a79f008 100644
--- a/lib/Parser.hs
+++ b/lib/Parser.hs
@@ -3,13 +3,12 @@
 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, satisfy, try)
+import Text.Megaparsec (Parsec, ShowErrorComponent (..), between, choice, chunk, errorBundlePretty, label, runParser, try)
 import Text.Megaparsec.Char
 import qualified Text.Megaparsec.Char.Lexer as L
 
@@ -42,12 +41,13 @@ reservedChars = "();:"
 
 pIdentifier :: Parser Text
 pIdentifier = try $ label "identifier" $ lexeme $ do
-    chars <- many $ satisfy isAllowed
-    let ident = T.pack chars
-    when (ident `elem` keywords) $ fail $ "Reserved word: " ++ T.unpack ident
+    firstChar <- letterChar <|> char '_'
+    rest <- many $ alphaNumChar <|> char '_'
+    let ident = T.pack (firstChar : rest)
+    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
@@ -127,12 +127,11 @@ pNumSort :: Parser IRExpr
 pNumSort = lexeme $ label "sort" $ eat "*" >> Level <$> L.decimal
 
 pSort :: Parser IRExpr
-pSort = try pNumSort <|> pStar
+pSort = lexeme $ try pNumSort <|> pStar
 
 pAxiom :: Parser IRDef
 pAxiom = lexeme $ label "axiom" $ do
-    skipSpace
-    eat "def"
+    eat "axiom"
     ident <- pIdentifier
     params <- pManyParams
     ascription <- fmap (flip (foldr (uncurry Pi)) params) pAscription
@@ -141,7 +140,6 @@ pAxiom = lexeme $ label "axiom" $ do
 
 pDef :: Parser IRDef
 pDef = lexeme $ label "definition" $ do
-    skipSpace
     eat "def"
     ident <- pIdentifier
     params <- pManyParams
@@ -152,13 +150,10 @@ pDef = lexeme $ label "definition" $ do
     pure $ Def ident ascription $ foldr (uncurry Abs) body params
 
 pIRDef :: Parser IRDef
-pIRDef = pDef <|> pAxiom
+pIRDef = pAxiom <|> pDef
 
 pTerm :: Parser IRExpr
-pTerm =
-    lexeme $
-        label "term" $
-            choice [between (char '(') (char ')') pIRExpr, pSort, pVar]
+pTerm = lexeme $ label "term" $ choice [pSort, pVar, between (char '(') (char ')') pIRExpr]
 
 pAppTerm :: Parser IRExpr
 pAppTerm = lexeme $ choice [pLAbs, pPAbs, pLet, pApp]
@@ -170,7 +165,7 @@ pAscription :: Parser IRExpr
 pAscription = lexeme $ try $ eat ":" >> label "type" pIRExpr
 
 pIRProgram :: Parser IRProgram
-pIRProgram = many pIRDef
+pIRProgram = skipSpace >> some pIRDef
 
 parserWrapper :: Parser a -> String -> Text -> Either String a
 parserWrapper p filename input = first errorBundlePretty $ runParser p filename input