perga/lib/Parser.hs

{-# LANGUAGE NamedFieldPuns #-}

module Parser (parseDef, parseDefEmpty, parseExpr, parseProgram, handleFile) where

import Check
import Control.Monad.Except
import Data.List (elemIndex, foldl, foldl1)
import qualified Data.Map.Strict as M
import qualified Data.Text as T
import Errors (Error (..))
import Eval
import Expr (Expr (..), incIndices)
import Preprocessor
import Relude.Extra.Lens
import Text.Megaparsec (ParsecT, ShowErrorComponent (..), between, choice, chunk, customFailure, errorBundlePretty, label, runParserT, try)
import Text.Megaparsec.Char
import qualified Text.Megaparsec.Char.Lexer as L

data InnerState = IS {_binders :: [Text], _env :: Env}

bindsL :: Lens' InnerState [Text]
bindsL = lens _binders setter
  where
    setter (IS{_env}) new = IS{_env, _binders = new}

envL :: Lens' InnerState Env
envL = lens _env setter
  where
    setter (IS{_binders}) new = IS{_env = new, _binders}

newtype TypeError = TE Error
    deriving (Eq, Ord)

type Parser = ParsecT TypeError Text (State InnerState)

instance ShowErrorComponent TypeError where
    showErrorComponent (TE e) = toString e

skipSpace :: Parser ()
skipSpace =
    L.space
        space1
        (L.skipLineComment "--")
        (L.skipBlockCommentNested "[*" "*]")

lexeme :: Parser a -> Parser a
lexeme = L.lexeme skipSpace

eat :: Text -> Parser ()
eat = void . lexeme . chunk

pIdentifier :: Parser Text
pIdentifier = label "identifier" $ lexeme $ do
    firstChar <- letterChar <|> char '_'
    rest <- many $ alphaNumChar <|> char '_'
    return $ T.pack (firstChar : rest)

pVar :: Parser Expr
pVar = label "variable" $ lexeme $ do
    name <- pIdentifier
    binders <- view bindsL <$> get
    pure (Var name . fromIntegral <$> elemIndex name binders ?: Free name)

defChoice :: NonEmpty Text -> Parser ()
defChoice options = lexeme $ label (T.unpack $ head options) $ void $ choice $ fmap chunk options

pParamGroup :: Parser [(Text, Expr)]
pParamGroup = lexeme $ label "parameter group" $ between (char '(') (char ')') $ do
    idents <- some pIdentifier
    eat ":"
    ty <- pExpr
    modify $ over bindsL $ flip (foldl $ flip (:)) idents
    pure $ zip idents (iterate incIndices ty)

pSomeParams :: Parser [(Text, Expr)]
pSomeParams = lexeme $ concat <$> some pParamGroup

pManyParams :: Parser [(Text, Expr)]
pManyParams = lexeme $ concat <$> many pParamGroup

withBinders :: Parser a -> Parser a
withBinders parser = do
    oldBinders <- view bindsL <$> get
    result <- parser
    modify $ set bindsL oldBinders
    pure result

pLAbs :: Parser Expr
pLAbs = lexeme $ label "λ-abstraction" $ withBinders $ do
    _ <- defChoice $ "λ" :| ["fun"]
    params <- pSomeParams
    _ <- defChoice $ "=>" :| ["⇒"]
    body <- pExpr
    pure $ foldr (uncurry Abs) body params

pPAbs :: Parser Expr
pPAbs = lexeme $ label "Π-abstraction" $ withBinders $ do
    _ <- defChoice $ "∏" :| ["forall", "∀"]
    params <- pSomeParams
    _ <- defChoice $ pure ","
    body <- pExpr
    pure $ foldr (uncurry Pi) body params

pArrow :: Parser Expr
pArrow = lexeme $ label "->" $ do
    a <- pAppTerm
    _ <- defChoice $ "->" :| ["→"]
    Pi "" a . incIndices <$> pExpr

pApp :: Parser Expr
pApp = lexeme $ foldl1 App <$> some pTerm

pStar :: Parser Expr
pStar = lexeme $ Star <$ eat "*"

pSquare :: Parser Expr
pSquare = lexeme $ Square <$ defChoice ("□" :| ["[]"])

checkAscription :: Text -> Expr -> Maybe Expr -> Parser ()
checkAscription ident value massert = do
    env <- get
    ty <- either (customFailure . TE) pure (checkType (env ^. envL) value)
    case massert of
        Nothing -> updateStateDefinition ident ty value
        Just assert -> do
            equiv <- either (customFailure . TE) pure (checkBeta (env ^. envL) ty assert)
            unless equiv $ customFailure $ TE $ NotEquivalent ty assert value
            updateStateDefinition ident assert value

updateStateDefinition :: Text -> Expr -> Expr -> Parser ()
updateStateDefinition ident ty value = do
    env <- get
    when (M.member ident (env ^. envL)) (customFailure $ TE $ DuplicateDefinition ident)
    modify $ over envL $ M.insert ident $ makeDef ty value

pAxiom :: Text -> Maybe Expr -> Parser ()
pAxiom ident Nothing = customFailure $ TE $ PNMissingType ident
pAxiom ident (Just ascription) = do
    eat "axiom"
    eat ";"
    updateStateDefinition ident ascription (Axiom ident)

pBody :: [(Text, Expr)] -> Text -> Maybe Expr -> Parser ()
pBody params ident ascription = do
    value <- flip (foldr (uncurry Abs)) params <$> pExpr
    checkAscription ident value ascription
    eat ";"

pDef :: Parser ()
pDef = lexeme $ label "definition" $ do
    skipSpace
    ident <- pIdentifier
    params <- pManyParams
    ascription <- fmap (flip (foldr (uncurry Pi)) params) <$> pAscription
    eat ":="
    choice [pAxiom ident ascription, pBody params ident ascription]

pTerm :: Parser Expr
pTerm =
    lexeme $
        label "term" $
            choice
                [ between (char '(') (char ')') pExpr
                , pStar
                , pSquare
                , pVar
                ]

pAppTerm :: Parser Expr
pAppTerm = lexeme $ pLAbs <|> pPAbs <|> pApp

pExpr :: Parser Expr
pExpr = lexeme $ try pArrow <|> pAppTerm

pAscription :: Parser (Maybe Expr)
pAscription = lexeme $ optional $ try $ defChoice (pure ":") >> label "type" pExpr

pProgram :: Parser Env
pProgram = lexeme $ skipSpace >> many pDef >> _env <$> get

emptyBinders :: Env -> InnerState
emptyBinders env = IS{_binders = [], _env = env}

parseDef :: Text -> State Env (Either String ())
parseDef input = do
    env <- get
    let (output, IS{_env}) = runState (runParserT pDef "" input) (emptyBinders env)
    put _env
    pure $ first errorBundlePretty output

parseExpr :: Env -> Text -> Either String Expr
parseExpr env input = first errorBundlePretty $ evalState (runParserT pExpr "" input) $ emptyBinders env

parseDefEmpty :: Env -> Text -> (Either String (), Env)
parseDefEmpty env input = runState (parseDef input) env

parseProgram :: Env -> Text -> Either String Env
parseProgram initial input = first errorBundlePretty $ evalState (runParserT pProgram "" input) $ emptyBinders initial

handleFile :: Env -> String -> ExceptT String IO Env
handleFile initial filename = do
    text <- toString `withExceptT` preprocess filename
    liftEither $ parseProgram initial text