| 
 Copyright    : ( c ) 2020 Fugue , Inc. 
 License      : Apache License , version 2.0 
 Maintainer   :  
 Stability    : experimental 
 Portability : POSIX 

Copyright   : (c) 2020 Fugue, Inc.
License     : Apache License, version 2.0
Maintainer  : 
Stability   : experimental
Portability : POSIX
-}
{-# LANGUAGE BangPatterns      #-}
{-# LANGUAGE OverloadedStrings #-}
# LANGUAGE RecordWildCards    #
# LANGUAGE TemplateHaskell    #
module Fregot.Parser.Sugar
    ( ParserOptions, poDefaultPackageName
    , defaultParserOptions

    , parseModule
    , rule
    , query
    , expr

    , Tok.string
    , scalar
    , array
    , objectOrSet
    ) where

import           Control.Applicative         ((<|>))
import           Control.Lens                ((^.))
import           Control.Lens.TH             (makeLenses)
import           Data.Either                 (partitionEithers)
import           Data.Functor                (($>))
import qualified Data.Scientific             as Scientific
import           Fregot.Parser.Internal
import qualified Fregot.Parser.Token         as Tok
import qualified Fregot.PrettyPrint          as PP
import           Fregot.Sources.SourceSpan
import           Fregot.Sugar
import           Prelude                     hiding (head)
import qualified Text.Parsec.Expr            as Parsec
import           Text.Parsec.Extended        ((<?>))
import qualified Text.Parsec.Extended        as Parsec
import qualified Text.Parsec.Indent.Explicit as Indent

data ParserOptions = ParserOptions
    { -- | A default package name which is used if there is no package name
      -- declaration in the module.  If this is not set, a package name
      -- declaration is required.
      _poDefaultPackageName :: Maybe PackageName
    } deriving (Show)

$(makeLenses ''ParserOptions)

defaultParserOptions :: ParserOptions
defaultParserOptions = ParserOptions
    { _poDefaultPackageName = Nothing
    }

parsePackageName :: FregotParser PackageName
parsePackageName =
    mkPackageName <$> Parsec.sepBy1 bit (Tok.symbol Tok.TPeriod)
  where
    bit = Tok.var <|>
        (Tok.symbol Tok.TIn    >> pure "in") <|>
        (Tok.symbol Tok.TEvery >> pure "every")

parseImportGut :: FregotParser ImportGut
parseImportGut = do
    pos     <- Parsec.getPosition
    pkgname <- parsePackageName
    case unPackageName pkgname of
        ("data"  : vs)  -> return $! ImportData   (mkPackageName vs)
        ("input" : vs)  -> return $! ImportInput  (mkPackageName vs)
        ("future" : vs) -> return $! ImportFuture (mkPackageName vs)
        _               -> Parsec.unexpectedAt pos $
            show (PP.pretty' pkgname) ++
            " (imports should start with `data.` or `input.`)"

parseModule :: ParserOptions -> FregotParser (Module SourceSpan Var)
parseModule po = withSourceSpan $ do
    _modulePackage <- parseModuleHead po
    _moduleImports <- Parsec.many parseModuleImport
    _modulePolicy  <- Parsec.many rule
    pure $ \_moduleAnn -> Module {..}

parseModuleHead :: ParserOptions -> FregotParser PackageName
parseModuleHead po = case po ^. poDefaultPackageName of
    Nothing  -> parser
    Just def -> Parsec.option def parser
  where
    parser = (<?> "'package' declaration") $
        Tok.symbol Tok.TPackage *> parsePackageName

parseModuleImport :: FregotParser (Import SourceSpan)
parseModuleImport = withSourceSpan $ do
    Tok.symbol Tok.TImport
    _importGut <- parseImportGut
    _importAs  <- Parsec.optionMaybe $ do
        Tok.symbol Tok.TAs
        var
    return $ \_importAnn -> Import {..}

var :: FregotParser Var
var = mkVar <$> Tok.var

rule :: FregotParser (Rule SourceSpan Var)
rule = Rule
    <$> parseRuleHead
    <*> Parsec.many parseRuleBody
    <*> Parsec.many parseRuleElse

parseRuleHead :: FregotParser (RuleHead SourceSpan Var)
parseRuleHead = withSourceSpan $ do
    _ruleDefault <- Parsec.option False $ Tok.symbol Tok.TDefault $> True
    _ruleName <- var
    _ruleArgs <- Parsec.optionMaybe $
        sepTrailing Tok.TLParen Tok.TRParen Tok.TComma term
    _ruleIndex <- Parsec.optionMaybe $ do
        Tok.symbol Tok.TLBracket
        t <- term
        expectToken Tok.TRBracket
        return t
    (_ruleAssign, _ruleValue) <- Parsec.choice
        [ Tok.symbol Tok.TAssign >> (,) True . Just <$> expr
        , Tok.symbol Tok.TUnify >> (,) False . Just <$> expr
        , pure (False, Nothing)
        ]
    return $ \_ruleAnn -> RuleHead {..}

parseRuleBody :: FregotParser (RuleBody SourceSpan Var)
parseRuleBody = do
    Tok.symbol Tok.TLBrace
    body <- query
    Tok.symbol Tok.TRBrace
    return body

parseRuleElse :: FregotParser (RuleElse SourceSpan Var)
parseRuleElse = withSourceSpan $ do
    Tok.symbol Tok.TElse
    _ruleElseValue <- Parsec.optionMaybe $ do
        Tok.symbol Tok.TUnify
        term
    _ruleElseBody <- parseRuleBody
    return $ \_ruleElseAnn -> RuleElse {..}

query :: FregotParser (RuleBody SourceSpan Var)
query = blockOrSemi ruleStatement

  | Parse either a block of lines , or lines separated by a semicolon , or both .
blockOrSemi :: FregotParser a -> FregotParser [a]
blockOrSemi linep =
    (do
        !pos <- Parsec.getPosition
        !l   <- linep
        (l :) <$> go pos) <|>
    (return [])
  where
    go pos0 =
        (do
            Tok.symbol Tok.TSemicolon
            !pos1 <- Parsec.getPosition
            !l    <- linep
            (l :) <$> go pos1) <|>
        (do
            !pos1 <- Parsec.getPosition
            if Parsec.sourceLine pos1 <= Parsec.sourceLine pos0
                then Parsec.parserFail "expected newline before next statement"
                else do
                    !l <- linep
                    (l :) <$> go pos1) <|>
        return []

ruleStatement :: FregotParser (RuleStatement SourceSpan Var)
ruleStatement =
    (withSourceSpan $ do
        (term1, term2) <- Parsec.try $ do
            Tok.symbol Tok.TSome
            term1 <- term
            term2 <- Parsec.optionMaybe $ do
                Tok.symbol Tok.TComma
                term
            Tok.symbol Tok.TIn
            pure (term1, term2)
        x <- expr
        let (k, v) = case term2 of
                Nothing -> (Nothing, term1)
                Just t2 -> (Just term1, t2)
        return $ \ann -> SomeInS ann k v x) <|>
    (withSourceSpan $ do
        (term1, term2) <- Parsec.try $ do
            Tok.symbol Tok.TEvery
            term1 <- term
            term2 <- Parsec.optionMaybe $ do
                Tok.symbol Tok.TComma
                term
            Tok.symbol Tok.TIn
            pure (term1, term2)
        x <- expr
        body <- parseRuleBody
        let (k, v) = case term2 of
                Nothing -> (Nothing, term1)
                Just t2 -> (Just term1, t2)
        return $ \ann -> EveryInS ann k v x body) <|>
    (withSourceSpan $ do
        Tok.symbol Tok.TSome
        vars <- Parsec.sepBy1 var (Tok.symbol Tok.TComma)
        return $ \ann -> SomeS ann vars) <|>
     (LiteralS <$> literal)

literal :: FregotParser (Literal SourceSpan Var)
literal = withSourceSpan $ do
    _literalNegation <- Parsec.option False $ Tok.symbol Tok.TNot $> True
    _literalExpr <- expr
    _literalWith <- Parsec.many parseWith
    return $ \_literalAnn -> Literal {..}

expr :: FregotParser (Expr SourceSpan Var)
expr = Parsec.buildExpressionParser
    [ [ binary Tok.TTimes  TimesO  Parsec.AssocLeft
      , binary Tok.TDivide DivideO Parsec.AssocLeft
      , binary Tok.TModulo ModuloO Parsec.AssocLeft
      ]

    , [ binary Tok.TPlus  PlusO  Parsec.AssocLeft
      , binary Tok.TMinus MinusO Parsec.AssocLeft
      ]

    , [ binary Tok.TLessThan           LessThanO           Parsec.AssocLeft
      , binary Tok.TLessThanOrEqual    LessThanOrEqualO    Parsec.AssocLeft
      , binary Tok.TGreaterThan        GreaterThanO        Parsec.AssocLeft
      , binary Tok.TGreaterThanOrEqual GreaterThanOrEqualO Parsec.AssocLeft
      ]

    , [ binary Tok.TBinAnd BinAndO Parsec.AssocLeft ]
    , [ binary Tok.TPipe   BinOrO  Parsec.AssocLeft ]

    , [ binary Tok.TEqual    UnifyO    Parsec.AssocLeft
      , binary Tok.TNotEqual NotEqualO Parsec.AssocLeft
      ]

    , [ binary Tok.TUnify  UnifyO  Parsec.AssocRight ]
    , [ binary Tok.TAssign AssignO Parsec.AssocRight ]
    ]
    expr1
  where
    expr1 = withSourceSpan $ do
        kv0 <- expr2
        (do
            kv1 <- Parsec.try $ do
                kv1 <- Parsec.optionMaybe $ do
                    Tok.symbol Tok.TComma
                    expr2
                Tok.symbol Tok.TIn
                pure kv1
            x <- expr2
            let (k, v) = case kv1 of
                    Nothing -> (Nothing, kv0)
                    Just v1 -> (Just kv0, v1)
            pure $ \ss -> InE ss k v x) <|>
            (pure $ \_ -> kv0)

    expr2 = withSourceSpan $
        (do
            Tok.symbol Tok.TLParen
            e <- expr
            expectToken Tok.TRParen
            return $ \ss -> ParensE ss e) <|>
        (do
            startI <- Indent.indentation
            t <- term
            ras <- Parsec.option [] $ Parsec.many1 $ Indent.same startI >> refArg
            case ras of
                [] -> pure $ \ss -> TermE ss t
                _ : _ -> pure $ \ss -> IndRefE ss t ras)

    binary tok op = Parsec.Infix (do
        Tok.symbol tok
        return $ \x y ->
            BinOpE (unsafeMergeSourceSpan (x ^. exprAnn) (y ^. exprAnn)) x op y)

term :: FregotParser (Term SourceSpan Var)
term = withSourceSpan $
    (do
        startI   <- Indent.indentation
        (v, vss) <- withSourceSpan $ var >>= \v -> return $ \vss -> (v, vss)
        refArgs  <- Parsec.option [] $ Parsec.many1 $ Indent.same startI >> refArg

        let isDotArg (RefDotArg _ _) = True
            isDotArg (RefBrackArg _) = False

        mbCallArgs <- case all isDotArg refArgs of
            False -> return Nothing
            True  -> Parsec.optionMaybe $
                sepTrailing Tok.TLParen Tok.TRParen Tok.TComma expr

        return $ \ss -> case (refArgs, mbCallArgs) of
            ([], Nothing)   -> VarT ss v
            (_,  Nothing)   -> RefT ss vss v refArgs
            (_, Just cargs) -> CallT ss
                (v : [v' | RefDotArg _ v' <- refArgs])
                cargs) <|>
    (do
        s <- scalar
        return $ \ss -> ScalarT ss s) <|>
    (do
        (x, body) <- comprehension Tok.TLBracket Tok.TRBracket term
        return $ \ss -> ArrayCompT ss x body) <|>
    (do
        a <- array expr
        return $ \ss -> ArrayT ss a) <|>
    (do
        (x, body) <- comprehension Tok.TLBrace Tok.TRBrace term
        return $ \ss -> SetCompT ss x body) <|>
    (do
        ((k, x), body) <- comprehension Tok.TLBrace Tok.TRBrace $ do
            k0 <- objectKey
            Tok.symbol Tok.TColon
            x0 <- term
            return (k0, x0)
        return $ \ss -> ObjectCompT ss k x body) <|>
    (do
        os <- objectOrSet objectKey expr
        return $ \ss -> case os of
            Left  o -> ObjectT ss o
            Right s -> SetT ss s)

refArg :: FregotParser (RefArg SourceSpan Var)
refArg =
    (do
        Tok.symbol Tok.TLBracket
        e <- expr
        Parsec.option () $ do
  Do n't call  if there is a comma because then we will
            -- really mess up the parse state here.
            Tok.symbol Tok.TComma
            Parsec.unexpected "token, expecting ]"
        expectToken Tok.TRBracket
        return $ RefBrackArg e) <|>
    (withSourceSpan $ do
        Tok.symbol Tok.TPeriod
        v <- var
        return $ \ss -> RefDotArg ss v)

scalar :: FregotParser Scalar
scalar =
    (String <$> Tok.string) <|>
    (do
        negation   <- Parsec.option False (True <$ Tok.symbol Tok.TMinus)
        intOrFloat <- Tok.intOrFloat
        return $! Number $! (if negation then negate else id) $
            case intOrFloat of
                Left  x -> fromIntegral x
                Right x -> Scientific.fromFloatDigits x) <|>
    (pure (Bool True) <* Tok.symbol Tok.TTrue) <|>
    (pure (Bool False) <* Tok.symbol Tok.TFalse) <|>
    (pure Null <* Tok.symbol Tok.TNull)

array :: FregotParser expr -> FregotParser [expr]
array = sepTrailing Tok.TLBracket Tok.TRBracket Tok.TComma

objectOrSet
    :: FregotParser key -> FregotParser value
    -> FregotParser (Either [(key, value)] [value])
objectOrSet key value = do
    items <- sepTrailing Tok.TLBrace Tok.TRBrace Tok.TComma item
    case partitionEithers items of
        (objitems, []) -> return $ Left objitems  -- `{}` is an empty object
        ([], setitems) -> return $ Right setitems
        _              -> Parsec.unexpected "mixed object and set"
  where
    item =
        (do
            !k <- Parsec.try $ key <* Tok.symbol Tok.TColon
            !e <- value
            return $ Left (k, e)) <|>
        (do
            !e <- value
            pure $ Right e)

objectKey :: FregotParser (ObjectKey SourceSpan Var)
objectKey = withSourceSpan $
    (do
        s <- scalar
        return $ \ss -> ScalarK ss s) <|>
    (do
        v       <- var
        refArgs <- Parsec.many refArg
        return $ \ss -> case refArgs of
            [] -> VarK ss v
            _  -> RefK ss v refArgs)

parseWith :: FregotParser (With SourceSpan Var)
parseWith = withSourceSpan $ do
    Tok.symbol Tok.TWith
    pos       <- Parsec.getPosition
    path      <- Parsec.sepBy1 var (Tok.symbol Tok.TPeriod)
    _withPath <- case path of
        "input" : p -> pure $! InputWithPath p
        "data"  : p -> pure $! DataWithPath  p
        k       : _ -> Parsec.unexpectedAt pos $
            show (PP.pretty' k) ++
            " (`with` statements should start with `data.` or `input.`)"
        -- Impossible because of the `Parsec.sepBy1` above.
        [] -> error "Internal error: empty with path"
    Tok.symbol Tok.TAs
    _withAs <- term
    return $ \_withAnn -> With {..}

-- | Parse a comprehension.  These use 'Parsec.try' and should be tried before
-- parsing regular arrays/objects etc.
comprehension
    :: Tok.Token       -- ^ '(' or '{'
    -> Tok.Token       -- ^ ')' or '}'
    -> FregotParser a  -- ^ Comprehension head
    -> FregotParser (a, RuleBody SourceSpan Var)
comprehension open close phead = do
    chead <- Parsec.try $ do
        Tok.symbol open
        x <- phead
        Tok.symbol Tok.TPipe
        return x
    cbody <- query
    Tok.symbol close
    return (chead, cbody)

-- | An expression between braces separated with commas.  There may be a
-- trailing comma.  We also need a start and end token.
sepTrailing
    :: Tok.Token         -- ^ '(' or '{'
    -> Tok.Token         -- ^ ')' or '}'
    -> Tok.Token         -- ^ ',' or ';'
  ^ 
    -> FregotParser [a]  -- ^ Result
sepTrailing open close comma p = do
    Tok.symbol open
    Parsec.choice
        [ Tok.symbol close >> return []
        , reverse <$> go []
        ]
  where
    go acc = do
        x <- p
        Parsec.choice
            [ Tok.symbol close >> return (x : acc)
            , do
                expectToken comma
                Parsec.choice
                    [ Tok.symbol close >> return (x : acc)
                    , go (x : acc)
                    ]
            ]

# LANGUAGE BangPatterns      #
# LANGUAGE OverloadedStrings #
 | A default package name which is used if there is no package name
 declaration in the module.  If this is not set, a package name
 declaration is required.
 really mess up the parse state here.
 `{}` is an empty object
 Impossible because of the `Parsec.sepBy1` above.
 | Parse a comprehension.  These use 'Parsec.try' and should be tried before
 parsing regular arrays/objects etc.
 ^ '(' or '{'
 ^ ')' or '}'
 ^ Comprehension head
 | An expression between braces separated with commas.  There may be a
 trailing comma.  We also need a start and end token.
 ^ '(' or '{'
 ^ ')' or '}'
 ^ ',' or ';'
 ^ Result
| 
 Copyright    : ( c ) 2020 Fugue , Inc. 
 License      : Apache License , version 2.0 
 Maintainer   :  
 Stability    : experimental 
 Portability : POSIX 

Copyright   : (c) 2020 Fugue, Inc.
License     : Apache License, version 2.0
Maintainer  : 
Stability   : experimental
Portability : POSIX
-}
# LANGUAGE RecordWildCards    #
# LANGUAGE TemplateHaskell    #
module Fregot.Parser.Sugar
    ( ParserOptions, poDefaultPackageName
    , defaultParserOptions

    , parseModule
    , rule
    , query
    , expr

    , Tok.string
    , scalar
    , array
    , objectOrSet
    ) where

import           Control.Applicative         ((<|>))
import           Control.Lens                ((^.))
import           Control.Lens.TH             (makeLenses)
import           Data.Either                 (partitionEithers)
import           Data.Functor                (($>))
import qualified Data.Scientific             as Scientific
import           Fregot.Parser.Internal
import qualified Fregot.Parser.Token         as Tok
import qualified Fregot.PrettyPrint          as PP
import           Fregot.Sources.SourceSpan
import           Fregot.Sugar
import           Prelude                     hiding (head)
import qualified Text.Parsec.Expr            as Parsec
import           Text.Parsec.Extended        ((<?>))
import qualified Text.Parsec.Extended        as Parsec
import qualified Text.Parsec.Indent.Explicit as Indent

data ParserOptions = ParserOptions
      _poDefaultPackageName :: Maybe PackageName
    } deriving (Show)

$(makeLenses ''ParserOptions)

defaultParserOptions :: ParserOptions
defaultParserOptions = ParserOptions
    { _poDefaultPackageName = Nothing
    }

parsePackageName :: FregotParser PackageName
parsePackageName =
    mkPackageName <$> Parsec.sepBy1 bit (Tok.symbol Tok.TPeriod)
  where
    bit = Tok.var <|>
        (Tok.symbol Tok.TIn    >> pure "in") <|>
        (Tok.symbol Tok.TEvery >> pure "every")

parseImportGut :: FregotParser ImportGut
parseImportGut = do
    pos     <- Parsec.getPosition
    pkgname <- parsePackageName
    case unPackageName pkgname of
        ("data"  : vs)  -> return $! ImportData   (mkPackageName vs)
        ("input" : vs)  -> return $! ImportInput  (mkPackageName vs)
        ("future" : vs) -> return $! ImportFuture (mkPackageName vs)
        _               -> Parsec.unexpectedAt pos $
            show (PP.pretty' pkgname) ++
            " (imports should start with `data.` or `input.`)"

parseModule :: ParserOptions -> FregotParser (Module SourceSpan Var)
parseModule po = withSourceSpan $ do
    _modulePackage <- parseModuleHead po
    _moduleImports <- Parsec.many parseModuleImport
    _modulePolicy  <- Parsec.many rule
    pure $ \_moduleAnn -> Module {..}

parseModuleHead :: ParserOptions -> FregotParser PackageName
parseModuleHead po = case po ^. poDefaultPackageName of
    Nothing  -> parser
    Just def -> Parsec.option def parser
  where
    parser = (<?> "'package' declaration") $
        Tok.symbol Tok.TPackage *> parsePackageName

parseModuleImport :: FregotParser (Import SourceSpan)
parseModuleImport = withSourceSpan $ do
    Tok.symbol Tok.TImport
    _importGut <- parseImportGut
    _importAs  <- Parsec.optionMaybe $ do
        Tok.symbol Tok.TAs
        var
    return $ \_importAnn -> Import {..}

var :: FregotParser Var
var = mkVar <$> Tok.var

rule :: FregotParser (Rule SourceSpan Var)
rule = Rule
    <$> parseRuleHead
    <*> Parsec.many parseRuleBody
    <*> Parsec.many parseRuleElse

parseRuleHead :: FregotParser (RuleHead SourceSpan Var)
parseRuleHead = withSourceSpan $ do
    _ruleDefault <- Parsec.option False $ Tok.symbol Tok.TDefault $> True
    _ruleName <- var
    _ruleArgs <- Parsec.optionMaybe $
        sepTrailing Tok.TLParen Tok.TRParen Tok.TComma term
    _ruleIndex <- Parsec.optionMaybe $ do
        Tok.symbol Tok.TLBracket
        t <- term
        expectToken Tok.TRBracket
        return t
    (_ruleAssign, _ruleValue) <- Parsec.choice
        [ Tok.symbol Tok.TAssign >> (,) True . Just <$> expr
        , Tok.symbol Tok.TUnify >> (,) False . Just <$> expr
        , pure (False, Nothing)
        ]
    return $ \_ruleAnn -> RuleHead {..}

parseRuleBody :: FregotParser (RuleBody SourceSpan Var)
parseRuleBody = do
    Tok.symbol Tok.TLBrace
    body <- query
    Tok.symbol Tok.TRBrace
    return body

parseRuleElse :: FregotParser (RuleElse SourceSpan Var)
parseRuleElse = withSourceSpan $ do
    Tok.symbol Tok.TElse
    _ruleElseValue <- Parsec.optionMaybe $ do
        Tok.symbol Tok.TUnify
        term
    _ruleElseBody <- parseRuleBody
    return $ \_ruleElseAnn -> RuleElse {..}

query :: FregotParser (RuleBody SourceSpan Var)
query = blockOrSemi ruleStatement

  | Parse either a block of lines , or lines separated by a semicolon , or both .
blockOrSemi :: FregotParser a -> FregotParser [a]
blockOrSemi linep =
    (do
        !pos <- Parsec.getPosition
        !l   <- linep
        (l :) <$> go pos) <|>
    (return [])
  where
    go pos0 =
        (do
            Tok.symbol Tok.TSemicolon
            !pos1 <- Parsec.getPosition
            !l    <- linep
            (l :) <$> go pos1) <|>
        (do
            !pos1 <- Parsec.getPosition
            if Parsec.sourceLine pos1 <= Parsec.sourceLine pos0
                then Parsec.parserFail "expected newline before next statement"
                else do
                    !l <- linep
                    (l :) <$> go pos1) <|>
        return []

ruleStatement :: FregotParser (RuleStatement SourceSpan Var)
ruleStatement =
    (withSourceSpan $ do
        (term1, term2) <- Parsec.try $ do
            Tok.symbol Tok.TSome
            term1 <- term
            term2 <- Parsec.optionMaybe $ do
                Tok.symbol Tok.TComma
                term
            Tok.symbol Tok.TIn
            pure (term1, term2)
        x <- expr
        let (k, v) = case term2 of
                Nothing -> (Nothing, term1)
                Just t2 -> (Just term1, t2)
        return $ \ann -> SomeInS ann k v x) <|>
    (withSourceSpan $ do
        (term1, term2) <- Parsec.try $ do
            Tok.symbol Tok.TEvery
            term1 <- term
            term2 <- Parsec.optionMaybe $ do
                Tok.symbol Tok.TComma
                term
            Tok.symbol Tok.TIn
            pure (term1, term2)
        x <- expr
        body <- parseRuleBody
        let (k, v) = case term2 of
                Nothing -> (Nothing, term1)
                Just t2 -> (Just term1, t2)
        return $ \ann -> EveryInS ann k v x body) <|>
    (withSourceSpan $ do
        Tok.symbol Tok.TSome
        vars <- Parsec.sepBy1 var (Tok.symbol Tok.TComma)
        return $ \ann -> SomeS ann vars) <|>
     (LiteralS <$> literal)

literal :: FregotParser (Literal SourceSpan Var)
literal = withSourceSpan $ do
    _literalNegation <- Parsec.option False $ Tok.symbol Tok.TNot $> True
    _literalExpr <- expr
    _literalWith <- Parsec.many parseWith
    return $ \_literalAnn -> Literal {..}

expr :: FregotParser (Expr SourceSpan Var)
expr = Parsec.buildExpressionParser
    [ [ binary Tok.TTimes  TimesO  Parsec.AssocLeft
      , binary Tok.TDivide DivideO Parsec.AssocLeft
      , binary Tok.TModulo ModuloO Parsec.AssocLeft
      ]

    , [ binary Tok.TPlus  PlusO  Parsec.AssocLeft
      , binary Tok.TMinus MinusO Parsec.AssocLeft
      ]

    , [ binary Tok.TLessThan           LessThanO           Parsec.AssocLeft
      , binary Tok.TLessThanOrEqual    LessThanOrEqualO    Parsec.AssocLeft
      , binary Tok.TGreaterThan        GreaterThanO        Parsec.AssocLeft
      , binary Tok.TGreaterThanOrEqual GreaterThanOrEqualO Parsec.AssocLeft
      ]

    , [ binary Tok.TBinAnd BinAndO Parsec.AssocLeft ]
    , [ binary Tok.TPipe   BinOrO  Parsec.AssocLeft ]

    , [ binary Tok.TEqual    UnifyO    Parsec.AssocLeft
      , binary Tok.TNotEqual NotEqualO Parsec.AssocLeft
      ]

    , [ binary Tok.TUnify  UnifyO  Parsec.AssocRight ]
    , [ binary Tok.TAssign AssignO Parsec.AssocRight ]
    ]
    expr1
  where
    expr1 = withSourceSpan $ do
        kv0 <- expr2
        (do
            kv1 <- Parsec.try $ do
                kv1 <- Parsec.optionMaybe $ do
                    Tok.symbol Tok.TComma
                    expr2
                Tok.symbol Tok.TIn
                pure kv1
            x <- expr2
            let (k, v) = case kv1 of
                    Nothing -> (Nothing, kv0)
                    Just v1 -> (Just kv0, v1)
            pure $ \ss -> InE ss k v x) <|>
            (pure $ \_ -> kv0)

    expr2 = withSourceSpan $
        (do
            Tok.symbol Tok.TLParen
            e <- expr
            expectToken Tok.TRParen
            return $ \ss -> ParensE ss e) <|>
        (do
            startI <- Indent.indentation
            t <- term
            ras <- Parsec.option [] $ Parsec.many1 $ Indent.same startI >> refArg
            case ras of
                [] -> pure $ \ss -> TermE ss t
                _ : _ -> pure $ \ss -> IndRefE ss t ras)

    binary tok op = Parsec.Infix (do
        Tok.symbol tok
        return $ \x y ->
            BinOpE (unsafeMergeSourceSpan (x ^. exprAnn) (y ^. exprAnn)) x op y)

term :: FregotParser (Term SourceSpan Var)
term = withSourceSpan $
    (do
        startI   <- Indent.indentation
        (v, vss) <- withSourceSpan $ var >>= \v -> return $ \vss -> (v, vss)
        refArgs  <- Parsec.option [] $ Parsec.many1 $ Indent.same startI >> refArg

        let isDotArg (RefDotArg _ _) = True
            isDotArg (RefBrackArg _) = False

        mbCallArgs <- case all isDotArg refArgs of
            False -> return Nothing
            True  -> Parsec.optionMaybe $
                sepTrailing Tok.TLParen Tok.TRParen Tok.TComma expr

        return $ \ss -> case (refArgs, mbCallArgs) of
            ([], Nothing)   -> VarT ss v
            (_,  Nothing)   -> RefT ss vss v refArgs
            (_, Just cargs) -> CallT ss
                (v : [v' | RefDotArg _ v' <- refArgs])
                cargs) <|>
    (do
        s <- scalar
        return $ \ss -> ScalarT ss s) <|>
    (do
        (x, body) <- comprehension Tok.TLBracket Tok.TRBracket term
        return $ \ss -> ArrayCompT ss x body) <|>
    (do
        a <- array expr
        return $ \ss -> ArrayT ss a) <|>
    (do
        (x, body) <- comprehension Tok.TLBrace Tok.TRBrace term
        return $ \ss -> SetCompT ss x body) <|>
    (do
        ((k, x), body) <- comprehension Tok.TLBrace Tok.TRBrace $ do
            k0 <- objectKey
            Tok.symbol Tok.TColon
            x0 <- term
            return (k0, x0)
        return $ \ss -> ObjectCompT ss k x body) <|>
    (do
        os <- objectOrSet objectKey expr
        return $ \ss -> case os of
            Left  o -> ObjectT ss o
            Right s -> SetT ss s)

refArg :: FregotParser (RefArg SourceSpan Var)
refArg =
    (do
        Tok.symbol Tok.TLBracket
        e <- expr
        Parsec.option () $ do
  Do n't call  if there is a comma because then we will
            Tok.symbol Tok.TComma
            Parsec.unexpected "token, expecting ]"
        expectToken Tok.TRBracket
        return $ RefBrackArg e) <|>
    (withSourceSpan $ do
        Tok.symbol Tok.TPeriod
        v <- var
        return $ \ss -> RefDotArg ss v)

scalar :: FregotParser Scalar
scalar =
    (String <$> Tok.string) <|>
    (do
        negation   <- Parsec.option False (True <$ Tok.symbol Tok.TMinus)
        intOrFloat <- Tok.intOrFloat
        return $! Number $! (if negation then negate else id) $
            case intOrFloat of
                Left  x -> fromIntegral x
                Right x -> Scientific.fromFloatDigits x) <|>
    (pure (Bool True) <* Tok.symbol Tok.TTrue) <|>
    (pure (Bool False) <* Tok.symbol Tok.TFalse) <|>
    (pure Null <* Tok.symbol Tok.TNull)

array :: FregotParser expr -> FregotParser [expr]
array = sepTrailing Tok.TLBracket Tok.TRBracket Tok.TComma

objectOrSet
    :: FregotParser key -> FregotParser value
    -> FregotParser (Either [(key, value)] [value])
objectOrSet key value = do
    items <- sepTrailing Tok.TLBrace Tok.TRBrace Tok.TComma item
    case partitionEithers items of
        ([], setitems) -> return $ Right setitems
        _              -> Parsec.unexpected "mixed object and set"
  where
    item =
        (do
            !k <- Parsec.try $ key <* Tok.symbol Tok.TColon
            !e <- value
            return $ Left (k, e)) <|>
        (do
            !e <- value
            pure $ Right e)

objectKey :: FregotParser (ObjectKey SourceSpan Var)
objectKey = withSourceSpan $
    (do
        s <- scalar
        return $ \ss -> ScalarK ss s) <|>
    (do
        v       <- var
        refArgs <- Parsec.many refArg
        return $ \ss -> case refArgs of
            [] -> VarK ss v
            _  -> RefK ss v refArgs)

parseWith :: FregotParser (With SourceSpan Var)
parseWith = withSourceSpan $ do
    Tok.symbol Tok.TWith
    pos       <- Parsec.getPosition
    path      <- Parsec.sepBy1 var (Tok.symbol Tok.TPeriod)
    _withPath <- case path of
        "input" : p -> pure $! InputWithPath p
        "data"  : p -> pure $! DataWithPath  p
        k       : _ -> Parsec.unexpectedAt pos $
            show (PP.pretty' k) ++
            " (`with` statements should start with `data.` or `input.`)"
        [] -> error "Internal error: empty with path"
    Tok.symbol Tok.TAs
    _withAs <- term
    return $ \_withAnn -> With {..}

comprehension
    -> FregotParser (a, RuleBody SourceSpan Var)
comprehension open close phead = do
    chead <- Parsec.try $ do
        Tok.symbol open
        x <- phead
        Tok.symbol Tok.TPipe
        return x
    cbody <- query
    Tok.symbol close
    return (chead, cbody)

sepTrailing
  ^ 
sepTrailing open close comma p = do
    Tok.symbol open
    Parsec.choice
        [ Tok.symbol close >> return []
        , reverse <$> go []
        ]
  where
    go acc = do
        x <- p
        Parsec.choice
            [ Tok.symbol close >> return (x : acc)
            , do
                expectToken comma
                Parsec.choice
                    [ Tok.symbol close >> return (x : acc)
                    , go (x : acc)
                    ]
            ]
