<commit_before>-- | A parser for comprehension language (CL) expressions.
--
-- @
-- t ::= [t] | (t, ..., t) | Int | Double | ...
--
-- e ::= table(n)::t
--     | (f1 e)::t
--     | (f2 e e)::t
--     | (op1 e)::t
--     | (e op2 e)::t
--     | (if e then e else e)::t
--     | l::t
--     | x::t
--     | [ e | qs ]::t
--     | (e, ..., e)::t
--     | (let x = e in e)::t
--
-- qs ::= x <- e, qs | e, qs | x <- e | e
-- @

module Database.DSH.CL.Parser where

import           Control.Applicative
import           Control.Monad
import           Data.List.NonEmpty       (NonEmpty ((:|)))
import qualified Data.List.NonEmpty       as N
import qualified Data.Text                as T

import           Text.Megaparsec
import qualified Text.Megaparsec.Lexer    as Lex

import           Database.DSH.CL.Lang
import qualified Database.DSH.Common.Lang as L
import qualified Database.DSH.Common.Type as T

type CLParser a = Parsec String a

--------------------------------------------------------------------------------
-- Lexer infrastructure

spaceConsumer :: CLParser ()
spaceConsumer = Lex.space (void spaceChar) empty empty

symbol :: String -> CLParser String
symbol = Lex.symbol spaceConsumer

lexeme :: CLParser a -> CLParser a
lexeme = Lex.lexeme spaceConsumer

parens :: CLParser a -> CLParser a
parens = between (symbol "(") (symbol ")")

brackets :: CLParser a -> CLParser a
brackets = between (symbol "[") (symbol "]")

braces :: CLParser a -> CLParser a
braces = between (symbol "{") (symbol "}")

colon :: CLParser ()
colon = void $ symbol ":"

pipe :: CLParser ()
pipe = void $ symbol "|"

comma :: CLParser ()
comma = void $ symbol ","

ident :: CLParser String
ident = lexeme (some alphaNumChar)

kw :: String -> CLParser ()
kw s = void $ lexeme (string s)

sep :: CLParser ()
sep = spaceConsumer

integerLit :: CLParser Int
integerLit = fromIntegral <$> lexeme Lex.integer

boolLit :: CLParser Bool
boolLit = kw "True" *> pure True <|> kw "False" *> pure False

stringLit :: CLParser String
stringLit = char '"' >> manyTill Lex.charLiteral (char '"')

doubleLit :: CLParser Double
doubleLit = lexeme Lex.float

unitLit :: CLParser ()
unitLit = void $ lexeme $ string "()"

--------------------------------------------------------------------------------
-- Parsing helpers

list :: CLParser a -> CLParser [a]
list p = brackets (p `sepBy` comma)

nonEmpty :: CLParser a -> CLParser (N.NonEmpty a)
nonEmpty p = brackets $ do
    x <- p
    void $ comma
    xs <- p `sepBy` comma
    return $ x :| xs

kwConst :: String -> a -> CLParser a
kwConst s c = kw s *> pure c



--------------------------------------------------------------------------------
-- Types

baseType :: CLParser ScalarType
baseType = try (kw "Int" *> pure T.IntT)
           <|> try (kw "Bool" *> pure T.BoolT)
           <|> try (kw "Double" *> pure T.DoubleT)
           <|> try (kw "()" *> pure T.UnitT)
           <|> try (kw "Decimal" *> pure T.DecimalT)
           <|> try (kw "Date" *> pure T.DateT)

typeExpr :: CLParser T.Type
typeExpr = T.ListT <$> brackets typeExpr
           <|> try (T.ScalarT <$> baseType)
           <|> T.TupleT <$> parens (typeExpr `sepBy1` comma)

typeAnnotation :: CLParser Type
typeAnnotation = colon >> colon >> typeExpr

--------------------------------------------------------------------------------
-- Table references

colName :: CLParser L.ColName
colName = L.ColName <$> ident

tableCols :: CLParser (N.NonEmpty L.Column)
tableCols = nonEmpty $ (,) <$> colName <*> baseType

tableKeys :: CLParser (N.NonEmpty L.Key)
tableKeys = nonEmpty $ L.Key <$> nonEmpty colName

baseTableSchema :: CLParser L.BaseTableSchema
baseTableSchema = do
    cols      <- tableCols
    keys      <- tableKeys
    return $ L.BaseTableSchema cols keys L.PossiblyEmpty

tableRef :: CLParser (Type -> Expr)
tableRef = do
    void $ string "table"
    (tabName, schema) <- parens ((,) <$> ident <*> (comma *> baseTableSchema))
    return $ \ty -> Table ty tabName schema

--------------------------------------------------------------------------------
-- Primitive functions:w

prim1 :: CLParser Prim1
prim1 =     try (kw "singleton" *> pure Singleton)
        <|> try (kw "only" *> pure Only)
        <|> try (kw "length" *> pure Length)
        <|> try (kw "concat" *> pure Concat)
        <|> try (kw "null" *> pure Null)
        <|> try (kw "sum" *> pure Sum)
        <|> try (kw "avg" *> pure Only)
        <|> try (kw "minimum" *> pure Minimum)
        <|> try (kw "maximum" *> pure Maximum)
        <|> try (kw "reverse" *> pure Reverse)
        <|> try (kw "and" *> pure And)
        <|> try (kw "or" *> pure Or)
        <|> try (kw "nub" *> pure Nub)
        <|> try (kw "number" *> pure Number)
        <|> try (kw "only" *> pure Only)
        <|> try (kw "sort" *> pure Sort)
        <|> try (kw "group" *> pure Group)

prim2 :: CLParser Prim2
prim2 =     try (kw "append" *> pure Append)
        <|> try (kw "zip" *> pure Zip)



--------------------------------------------------------------------------------
-- Literals

baseLit :: CLParser L.ScalarVal
baseLit =     try (L.DoubleV <$> doubleLit)
          <|> try (L.IntV <$> integerLit)
          <|> try (L.BoolV <$> boolLit)
          <|> try (L.StringV <$> T.pack <$> stringLit)
          <|> try (unitLit *> pure L.UnitV)

literal :: CLParser L.Val
literal = L.ScalarV <$> baseLit

--------------------------------------------------------------------------------
-- Infix operators
binNumOp :: CLParser L.BinNumOp
binNumOp =     try (kwConst "+" L.Add)
           <|> try (kwConst "-" L.Sub)
           <|> try (kwConst "/" L.Div)
           <|> try (kwConst "*" L.Mul)
           <|> try (kwConst "%" L.Mod)

binRelOp :: CLParser L.BinRelOp<commit_msg>Add literal lists and tuples to CL parser.<commit_after>-- | A parser for comprehension language (CL) expressions.
--
-- @
-- t ::= [t] | (t, ..., t) | Int | Double | ...
--
-- e ::= table(n)::t
--     | (prim1 e)::t
--     | (prim2 e e)::t
--     | (prefixOp e)::t
--     | (e infixOp e)::t
--     | (if e then e else e)::t
--     | l::t
--     | x::t
--     | [ e | qs ]::t
--     | (e, ..., e)::t
--     | (let x = e in e)::t
--
-- qs ::= x <- e, qs | e, qs | x <- e | e
-- @

module Database.DSH.CL.Parser
    ( parseCL
    ) where

import           Control.Applicative
import           Control.Monad
import           Data.List.NonEmpty       (NonEmpty ((:|)))
import qualified Data.List.NonEmpty       as N
import qualified Data.Text                as T

import           Text.Megaparsec
import qualified Text.Megaparsec.Lexer    as Lex

import           Database.DSH.CL.Lang
import qualified Database.DSH.Common.Lang as L
import qualified Database.DSH.Common.Type as T

type CLParser a = Parsec String a

--------------------------------------------------------------------------------
-- Lexer infrastructure

spaceConsumer :: CLParser ()
spaceConsumer = Lex.space (void spaceChar) empty empty

symbol :: String -> CLParser String
symbol = Lex.symbol spaceConsumer

lexeme :: CLParser a -> CLParser a
lexeme = Lex.lexeme spaceConsumer

parens :: CLParser a -> CLParser a
parens = between (symbol "(") (symbol ")")

brackets :: CLParser a -> CLParser a
brackets = between (symbol "[") (symbol "]")

braces :: CLParser a -> CLParser a
braces = between (symbol "{") (symbol "}")

colon :: CLParser ()
colon = void $ symbol ":"

pipe :: CLParser ()
pipe = void $ symbol "|"

comma :: CLParser ()
comma = void $ symbol ","

ident :: CLParser String
ident = lexeme (some alphaNumChar)

kw :: String -> CLParser ()
kw s = void $ lexeme (string s)

sep :: CLParser ()
sep = spaceConsumer

integerLit :: CLParser Int
integerLit = fromIntegral <$> lexeme Lex.integer

boolLit :: CLParser Bool
boolLit = kw "True" *> pure True <|> kw "False" *> pure False

stringLit :: CLParser String
stringLit = char '"' >> manyTill Lex.charLiteral (char '"')

doubleLit :: CLParser Double
doubleLit = lexeme Lex.float

unitLit :: CLParser ()
unitLit = void $ lexeme $ string "()"

--------------------------------------------------------------------------------
-- Parsing helpers

list :: CLParser a -> CLParser [a]
list p = brackets (p `sepBy` comma)

nonEmpty :: CLParser a -> CLParser (N.NonEmpty a)
nonEmpty p = brackets $ do
    x <- p
    void $ comma
    xs <- p `sepBy` comma
    return $ x :| xs

tuple :: CLParser a -> CLParser [a]
tuple p = parens (p `sepBy1` comma)

kwConst :: String -> a -> CLParser a
kwConst s c = kw s *> pure c

--------------------------------------------------------------------------------
-- Types

baseType :: CLParser ScalarType
baseType =     try (kw "Int" *> pure T.IntT)
           <|> try (kw "Bool" *> pure T.BoolT)
           <|> try (kw "Double" *> pure T.DoubleT)
           <|> try (kw "()" *> pure T.UnitT)
           <|> try (kw "Decimal" *> pure T.DecimalT)
           <|> try (kw "Date" *> pure T.DateT)

typeExpr :: CLParser T.Type
typeExpr = T.ListT <$> brackets typeExpr
           <|> try (T.ScalarT <$> baseType)
           <|> T.TupleT <$> tuple typeExpr

typeAnnotation :: CLParser Type
typeAnnotation = colon >> colon >> typeExpr

--------------------------------------------------------------------------------
-- Table references

colName :: CLParser L.ColName
colName = L.ColName <$> ident

tableCols :: CLParser (N.NonEmpty L.Column)
tableCols = nonEmpty $ (,) <$> colName <*> baseType

tableKeys :: CLParser (N.NonEmpty L.Key)
tableKeys = nonEmpty $ L.Key <$> nonEmpty colName

baseTableSchema :: CLParser L.BaseTableSchema
baseTableSchema = do
    cols      <- tableCols
    keys      <- tableKeys
    return $ L.BaseTableSchema cols keys L.PossiblyEmpty

tableRef :: CLParser (Type -> Expr)
tableRef = do
    void $ string "table"
    (tabName, schema) <- parens ((,) <$> ident <*> (comma *> baseTableSchema))
    return $ \ty -> Table ty tabName schema

--------------------------------------------------------------------------------
-- Primitive functions:w

prim1 :: CLParser Prim1
prim1 =     try (kw "singleton" *> pure Singleton)
        <|> try (kw "only" *> pure Only)
        <|> try (kw "length" *> pure Length)
        <|> try (kw "concat" *> pure Concat)
        <|> try (kw "null" *> pure Null)
        <|> try (kw "sum" *> pure Sum)
        <|> try (kw "avg" *> pure Only)
        <|> try (kw "minimum" *> pure Minimum)
        <|> try (kw "maximum" *> pure Maximum)
        <|> try (kw "reverse" *> pure Reverse)
        <|> try (kw "and" *> pure And)
        <|> try (kw "or" *> pure Or)
        <|> try (kw "nub" *> pure Nub)
        <|> try (kw "number" *> pure Number)
        <|> try (kw "only" *> pure Only)
        <|> try (kw "sort" *> pure Sort)
        <|> try (kw "group" *> pure Group)

prim2 :: CLParser Prim2
prim2 =     try (kw "append" *> pure Append)
        <|> try (kw "zip" *> pure Zip)



--------------------------------------------------------------------------------
-- Literals

baseLit :: CLParser L.ScalarVal
baseLit =     try (L.DoubleV <$> doubleLit)
          <|> try (L.IntV <$> integerLit)
          <|> try (L.BoolV <$> boolLit)
          <|> try (L.StringV <$> T.pack <$> stringLit)
          <|> try (unitLit *> pure L.UnitV)

literal :: CLParser L.Val
literal =     try (L.ScalarV <$> baseLit)
          <|> try (L.ListV <$> list literal)
          <|> try (L.TupleV <$> tuple literal)

--------------------------------------------------------------------------------
-- Infix operators
binNumOp :: CLParser L.BinNumOp
binNumOp =     try (kwConst "+" L.Add)
           <|> try (kwConst "-" L.Sub)
           <|> try (kwConst "/" L.Div)
           <|> try (kwConst "*" L.Mul)
           <|> try (kwConst "%" L.Mod)

binRelOp :: CLParser L.BinRelOp
{'fraction_non_alphanumeric': 0.2425280465873084, 'fraction_numerical': 0.001712768690588336, 'mean_word_length': 1.789109147360879, 'pattern_counts': {'":': 2, '<': 119, '<?xml version=': 0, '>': 202, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1757520', 'n_tokens_mistral': 3846, 'n_tokens_neox': 3651, 'n_words': 1111}