geolog-lsp/datalog/src/Datalog/Parser.hs

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE MultilineStrings #-}
{- HLINT ignore "Use void" -}

module Datalog.Parser (
    parseTerm,
    parseRule,
    parseProgram,
    SrcLoc (..),
)
where

import Data.Text (Text)
import Data.Text qualified as T
import Data.Void (Void)
import Datalog.Syntax (Atom' (..), Program' (..), Rule' (..))
import Datalog.Syntax hiding (Atom, Program, Rule, Term)
import Text.Megaparsec
import Text.Megaparsec.Char
import Text.Megaparsec.Char.Lexer qualified as L

type Parser = Parsec Void Text

data SrcLoc = SrcLoc
    { startLine :: Int
    , startCol :: Int
    , endLine :: Int
    , endCol :: Int
    }
    deriving (Show)

type Atom = Atom' SrcLoc
type Term = Term' SrcLoc
type Rule = Rule' SrcLoc
type Program = Program' SrcLoc

lexeme :: (MonadParsec e Text m) => m a -> m a
lexeme = L.lexeme whitespace

symbol :: (MonadParsec e Text m) => Text -> m Text
symbol = L.symbol whitespace

whitespace :: (MonadParsec e Text m) => m ()
whitespace =
    L.space
        space1
        (L.skipLineComment "--")
        (L.skipBlockComment "{-" "-}")

parens :: (MonadParsec e Text m) => m a -> m a
parens = between (symbol "(") (symbol ")")

comma, period :: (MonadParsec e Text m) => m ()
comma = () <$ symbol ","
period = () <$ symbol "."

parseCon :: Parser Term
parseCon = parseThingWithSub Con $ ConId . T.pack <$> lexeme (liftA2 (:) (numberChar <|> lowerChar) (many alphaNumChar))

parseVar :: Parser Term
parseVar = parseThingWithSub Var $ VarId . T.pack <$> lexeme (liftA2 (:) upperChar (many alphaNumChar))

parseTerm :: Parser Term
parseTerm = parseVar <|> parseCon

parseAtom :: Parser Atom
parseAtom = parseThingWithSub (\loc (rel, args) -> Atom loc rel args) do
    rel <- RelId . T.pack <$> lexeme (liftA2 (:) lowerChar (many alphaNumChar))
    args <- parens (parseTerm `sepBy` comma)
    return (rel, args)

parseQuery :: Parser [Atom]
parseQuery = parseAtom `sepBy` comma

parseRule :: Parser Rule
parseRule =
    parseThingWithSub (\loc (a, as) -> Rule loc a as) $
        try rule1 <|> rule2
  where
    rule1 = do
        headAtom <- parseAtom
        period
        return (headAtom, [])
    rule2 = do
        headAtom <- parseAtom <* symbol ":-"
        bodyAtoms <- parseQuery
        period
        return (headAtom, bodyAtoms)

parseProgram :: Parser Program
parseProgram = parseThingWithSub Program (many parseRule)

parseThingWithSub :: (SrcLoc -> c -> f SrcLoc) -> Parser c -> Parser (f SrcLoc)
parseThingWithSub f parseSub = do
    SourcePos _ sl sc <- getSourcePos
    c <- parseSub
    SourcePos _ el ec <- getSourcePos
    pure $ f (SrcLoc (unPos sl) (unPos sc) (unPos el) (unPos ec)) c