FunctionStatement.hs

module Compiler.AST.FunctionStatement where

import Control.Applicative
import Data.List.NonEmpty as NE
import Data.Text as T
import Data.Maybe (isJust)

import Common
import Compiler.AST.Common
import Compiler.AST.Expression
import Compiler.AST.Parser.Common
import Compiler.Lexer
import Compiler.Lexer.Comments
import Parser.Lib
import Parser.Parser
import Test.Common

data FunctionStatementWithLoc
  = FunctionStatementWithLoc FunctionStatement Location
  deriving Show

instance Eq FunctionStatementWithLoc where
  (FunctionStatementWithLoc fs1 _) == (FunctionStatementWithLoc fs2 _) = fs1 == fs2

instance ToSource FunctionStatementWithLoc where
  toSourcePretty i (FunctionStatementWithLoc fs _) = toSourcePretty i fs

instance HasAstParser FunctionStatementWithLoc where
  astParser = nameParser "Statement" $ do
    loc <- getParserLocation
    fsr <- astParser
    pure $ FunctionStatementWithLoc fsr loc

instance HasGen FunctionStatementWithLoc where
  getGen = FunctionStatementWithLoc <$> getGen <*> (pure emptyLocation)

type IsGlobal = Bool

data FunctionStatement
  = Let Subscript IsGlobal ExpressionWithLoc
  | Call ExpressionWithLoc
  | If ExpressionWithLoc (NonEmpty FunctionStatementWithLoc) (NonEmpty FunctionStatementWithLoc)
  | MultiIf ExpressionWithLoc (NonEmpty FunctionStatementWithLoc) (NonEmpty (ExpressionWithLoc, NonEmpty FunctionStatementWithLoc)) (Maybe (NonEmpty FunctionStatementWithLoc))
  | IfThen ExpressionWithLoc (NonEmpty FunctionStatementWithLoc)
  | For Identifier ExpressionWithLoc ExpressionWithLoc (NonEmpty FunctionStatementWithLoc) (Maybe ExpressionWithLoc)
  | ForEach Identifier ExpressionWithLoc (NonEmpty FunctionStatementWithLoc)
  | While ExpressionWithLoc (NonEmpty FunctionStatementWithLoc)
  | Loop (NonEmpty FunctionStatementWithLoc)
  | Return ExpressionWithLoc
  | Yield ExpressionWithLoc
  | Break
  | Pass
  | FnComment Comment
  deriving (Show, Eq)

instance ToSource FunctionStatement where
  toSourcePretty i Break = T.concat [indent i, toSource KwBreak]
  toSourcePretty i Pass = T.concat [indent i, toSource KwPass]
  toSourcePretty i (FnComment c) = stripEnd (toSourcePretty i c)
  toSourcePretty i (Let idf False expr) =
    T.concat [indent i, toSource KwLet, wst, toSource idf, wst, toSource KwAssignment, wst, toSource expr]
  toSourcePretty i (Let idf True expr) =
    T.concat [indent i, toSource KwLet, wst, toSource KwGlobal, wst, toSource idf, wst, toSource KwAssignment, wst, toSource expr]
  toSourcePretty i (Call expr) =
    T.concat $ [indent i, toSource expr]
  toSourcePretty i (Return expr) = T.concat $ [indent i, toSource KwReturn, wst, toSource expr]
  toSourcePretty i (Yield expr) = T.concat $ [indent i, toSource KwYield, wst, toSource expr]
  toSourcePretty i (ForEach idf expr1 stms) =
    T.concat $
      [indent i, toSource KwForEach, wst, toSource expr1, wst, toSource KwAs, wst, toSource idf, wst, nlt] <>
      [toSourcePretty (i+1) (NE.toList stms)] <>
      [nlt, indent i, toSource KwEndForEach]
  toSourcePretty i (For idf expr1 expr2 stms Nothing) =
    T.concat $
      [indent i, toSource KwFor, wst, toSource idf, wst, toSource KwAssignment, wst, toSource expr1, wst, toSource KwTo, wst, toSource expr2, nlt] <>
      [toSourcePretty (i+1) (NE.toList stms)] <>
      [nlt, indent i, toSource KwEndFor]
  toSourcePretty i (For idf expr1 expr2 stms (Just stepExpr)) =
    T.concat $
      [indent i, toSource KwFor, wst, toSource idf, wst, toSource KwAssignment, wst, toSource expr1, wst, toSource KwTo, wst, toSource expr2, wst, toSource KwStep, wst, toSource stepExpr, nlt] <>
      [toSourcePretty (i+1) (NE.toList stms)] <>
      [nlt, indent i, toSource KwEndFor]
  toSourcePretty i (IfThen expr stms) =
    T.concat $
      [indent i, toSource KwIf, wst, toSource expr, wst, toSource KwThen, nlt] <>
      [toSourcePretty (i+1) (NE.toList stms)] <>
      [nlt, indent i, toSourcePretty i KwEndIf]
  toSourcePretty i (MultiIf expr stms1 rst mstms2) =
    T.concat $
      [indent i, toSource KwIf, wst, toSource expr, wst, toSource KwThen, nlt] <>
      [toSourcePretty (i+1) (NE.toList stms1)] <>
      (toSourceElseIf <$> (NE.toList rst)) <>
      [nlt, toSourceElse] <>
      [indent i, toSource KwEndIf]
    where
      toSourceElse :: Text
      toSourceElse = case mstms2 of
        Just stms -> T.concat $
          [indent i, toSource KwElse, nlt] <> [toSourcePretty (i+1) (NE.toList stms)] <> [nlt]
        Nothing -> ""
      toSourceElseIf :: (ExpressionWithLoc, NonEmpty FunctionStatementWithLoc) -> Text
      toSourceElseIf (expr1, stms3) =
        T.concat $
          [nlt, indent i, toSource KwElseIf, wst, toSource expr1, wst, toSource KwThen, nlt] <>
          [toSourcePretty (i+1) (NE.toList stms3)]
  toSourcePretty i (If expr stms1 stms2) =
    T.concat $
      [indent i, toSource KwIf, wst, toSource expr, wst, toSource KwThen, nlt] <>
      [toSourcePretty (i+1) (NE.toList stms1)] <>
      [nlt, indent i, toSource KwElse, nlt] <>
      [toSourcePretty (i+1) (NE.toList stms2)] <>
      [nlt, indent i, toSource KwEndIf]
  toSourcePretty i (While expr1 stms) =
    T.concat $
      [indent i, toSource KwWhile, wst, toSource expr1, nlt] <>
      [toSourcePretty (i+1) (NE.toList stms)] <>
      [nlt, indent i, toSource KwEndWhile]
  toSourcePretty i (Loop stms) =
    T.concat $
      [indent i, toSource KwLoop, nlt] <>
      [toSourcePretty (i+1) (NE.toList stms)] <>
      [nlt, indent i, toSource KwEndLoop]

instance HasGen FunctionStatement where
  getGen = recursive choice
    [ Let <$> getGen <*> bool <*> getGen
    , Call <$> getGen
    , Return <$> getGen
    , Yield <$> getGen
    , FnComment <$> getGen
    , pure Break
    , pure Pass
    ]
    [ If <$> getGen <*> (nonEmptyGen getGen) <*> (nonEmptyGen getGen)
    , IfThen <$> getGen <*> (nonEmptyGen getGen)
    , MultiIf <$> getGen <*> (nonEmptyGen getGen) <*> (nonEmptyGen ((,) <$> getGen <*> (nonEmptyGen getGen))) <*> (Test.Common.maybe (nonEmptyGen getGen))
    , For <$> getGen <*> getGen <*> getGen <*> (nonEmptyGen getGen) <*> getGen
    , Loop <$> (nonEmptyGen getGen)
    , While <$> getGen <*> (nonEmptyGen getGen)
    ]

instance HasAstParser FunctionStatement where
  astParser = nameParser "Function statement (raw)" $ do
    (whitespaceNL <|> whitespace)
    functionStatementParser

functionStatementParser :: AstParser FunctionStatement
functionStatementParser
    = ifParser
    <|> letParser
    <|> forEachParser
    <|> forParser
    <|> whileParser
    <|> loopParser
    <|> returnParser
    <|> yieldParser
    <|> breakParser
    <|> passParser
    <|> callStatementParser
    <|> (FnComment <$> (surroundWs parseComment))

letParser :: AstParser FunctionStatement
letParser = nameParser "Let statement" $ do
  surroundWs_ (parseKeyword KwLet)
  (isJust -> isGlobal) <- surroundWs (optional (parseKeyword KwGlobal))
  idf <- surroundWs (mandatory parseSubscript)
  surroundWs_ (mandatory (parseToken "assignment" isAssignment))
  expr <- surroundWs (mandatory (astParser @ExpressionWithLoc))
  pure $ Let idf isGlobal expr
  where
    isAssignment (TkKeyword KwAssignment) = Just ()
    isAssignment _                        = Nothing

callStatementParser :: AstParser FunctionStatement
callStatementParser = nameParser "Function call" $ do
  expr <- surroundWs (astParser @ExpressionWithLoc)
  pure $ Call expr

loopParser :: AstParser FunctionStatement
loopParser = nameParser "Loop Statement" $ do
  surroundWs_ (parseKeyword KwLoop)
  stms <- mandatory (some $ astParser @FunctionStatementWithLoc)
  surroundWs_ (mandatory (parseKeyword KwEndLoop))
  pure $  Loop (NE.fromList stms)

whileParser :: AstParser FunctionStatement
whileParser = nameParser "While Statement" $ do
  surroundWs_ (parseKeyword KwWhile)
  guardBool <- surroundWs (mandatory $ astParser @ExpressionWithLoc)
  stms <- mandatory (some $ astParser @FunctionStatementWithLoc)
  surroundWs_ (mandatory (parseKeyword KwEndWhile))
  pure $  While guardBool (NE.fromList stms)

forParser :: AstParser FunctionStatement
forParser = nameParser "For Statement" $ do
  surroundWs_ (parseKeyword KwFor)
  idf <- surroundWs (mandatory parseIdentifier)
  surroundWs_ (mandatory $ parseKeyword KwAssignment)
  startExp <- surroundWs (mandatory $ astParser @ExpressionWithLoc)
  surroundWs_ (mandatory (parseKeyword KwTo))
  endExp <- surroundWs (mandatory $ astParser @ExpressionWithLoc)
  mStepExpr <- surroundWs (optional (parseKeyword KwStep)) >>= \case
    Just _ -> Just <$> surroundWs (mandatory $ astParser @ExpressionWithLoc)
    Nothing -> pure Nothing
  stms <- mandatory (some $ astParser @FunctionStatementWithLoc)
  surroundWs_ (mandatory (parseKeyword KwEndFor))
  pure $ For idf startExp endExp (NE.fromList stms) mStepExpr

forEachParser :: AstParser FunctionStatement
forEachParser = nameParser "ForEach Statement" $ do
  surroundWs_ (parseKeyword KwForEach)
  expr <- surroundWs (mandatory $ astParser @ExpressionWithLoc)
  surroundWs_ (mandatory $ parseKeyword KwAs)
  idf <- surroundWs (mandatory parseIdentifier)
  stms <- mandatory (some $ astParser @FunctionStatementWithLoc)
  surroundWs_ (mandatory (parseKeyword KwEndForEach))
  pure $ ForEach idf expr (NE.fromList stms)

ifParser :: AstParser FunctionStatement
ifParser = nameParser "If Statement" $ do
  surroundWs_ (parseKeyword KwIf)
  expr <- surroundWs (mandatory $ astParser @ExpressionWithLoc)
  surroundWs_ (mandatory $ parseKeyword KwThen)
  stms1 <- mandatory (some $ astParser @FunctionStatementWithLoc)
  surroundWs (mandatory (parseKeyword KwElseIf <|> parseKeyword KwElse <|> parseKeyword KwEndIf)) >>= \case
    KwElseIf -> do
      expr' <- surroundWs (mandatory $ astParser @ExpressionWithLoc)
      surroundWs_ (mandatory $ parseKeyword KwThen)
      stms2 <- mandatory (some $ astParser @FunctionStatementWithLoc)
      remaining <- many $ do
        surroundWs_ $ parseKeyword KwElseIf
        expr1 <- surroundWs (mandatory $ astParser @ExpressionWithLoc)
        surroundWs_ (mandatory $ parseKeyword KwThen)
        stms3 <- mandatory (some $ astParser @FunctionStatementWithLoc)
        pure (expr1, NE.fromList stms3)
      surroundWs (mandatory (parseKeyword KwElse <|> parseKeyword KwEndIf)) >>= \case
        KwElse ->  do
          elseStms <- mandatory (some $ astParser @FunctionStatementWithLoc)
          surroundWs_ (mandatory $ parseKeyword KwEndIf)
          pure $ MultiIf expr (NE.fromList stms1) ((expr', NE.fromList stms2) :| remaining) (Just $ NE.fromList elseStms)
        KwEndIf ->
          pure $ MultiIf expr (NE.fromList stms1) ((expr', NE.fromList stms2) :| remaining) Nothing
        _ -> fail "Impossible"
    KwElse -> do
      stms2 <- mandatory (some $ astParser @FunctionStatementWithLoc)
      surroundWs_ (mandatory $ parseKeyword KwEndIf)
      pure $ If expr (NE.fromList stms1) (NE.fromList stms2)
    KwEndIf ->
      pure $ IfThen expr (NE.fromList stms1)
    _ -> fail "Impossible"

returnParser :: AstParser FunctionStatement
returnParser = nameParser "return statement" $ do
  surroundWs_ (parseKeyword KwReturn)
  Return <$> (surroundWs $ astParser @ExpressionWithLoc)

yieldParser :: AstParser FunctionStatement
yieldParser = nameParser "yield statement" $ do
  surroundWs_ (parseKeyword KwYield)
  setGeneratorFlag True
  Yield <$> (surroundWs $ astParser @ExpressionWithLoc)

breakParser :: AstParser FunctionStatement
breakParser = nameParser "break statement" $ do
  surroundWs_ (parseKeyword KwBreak)
  pure Break

passParser :: AstParser FunctionStatement
passParser = nameParser "pass statement" $ do
  surroundWs_ (parseKeyword KwPass)
  pure Pass