FunctionDef.hs

module Compiler.AST.FunctionDef where

import Common
import Control.DeepSeq
import GHC.Generics
import Control.Applicative
import Data.List as L
import Data.List.NonEmpty as NE
import Data.Maybe
import Data.Text as T

import Compiler.AST.Common
import Compiler.AST.FunctionStatement
import Compiler.AST.Parser.Common
import Compiler.Lexer
import Parser.Lib
import Test.Common

data FunctionDef
  = FunctionDef Bool Identifier [Identifier] (NonEmpty FunctionStatementWithLoc)
  deriving (Show, Generic, Eq)

instance NFData FunctionDef

instance HasAstParser FunctionDef where
  astParser = nameParser "Function Definition" $ do
    surroundWs_ (parseKeyword KwProc)
    fnName <- surroundWs parseIdentifier
    args <- parseItemListInParen parseIdentifier
    setGeneratorFlag False
    stms <- mandatory (some (surroundWs (astParser @FunctionStatementWithLoc)))
    mandatory $ surroundWs_ $ parseKeyword KwEndProc
    isGenerator <- getGeneratorFlag
    setGeneratorFlag False
    pure $ FunctionDef isGenerator fnName (fromMaybe [] (NE.toList <$> args)) (NE.fromList stms)

instance ToSource FunctionDef where
  toSourcePretty i (FunctionDef _ name args stms) =
    T.concat $
      [nlt, indent i, toSource KwProc, wst, toSource name, toSource DlParenOpen] <>
      (L.intersperse (toSource DlComma <> " ") (toSource <$> args)) <>
      [toSource DlParenClose] <> [nlt] <>
      [toSourcePretty (i+1) (NE.toList stms)] <> [nlt, indent i, toSource KwEndProc]

instance HasGen FunctionDef where
  getGen = do
    isGen <- bool
    stms <- case isGen of
      True -> functionStmsWithYield
      False -> functionStmsWithoutYield
    (FunctionDef isGen) <$> getGen <*> getGen <*> (pure stms)

hasYield :: NonEmpty FunctionStatementWithLoc -> Bool
hasYield fustms = case fustms of
  ((FunctionStatementWithLoc stm _) :|  []) -> hasYield' stm
  ((FunctionStatementWithLoc stm _) :|  (h:t)) -> if hasYield' stm
    then True
    else hasYield (h :| t)
  where
    hasYield' = \case
      Let _ _ _ -> False
      Call _ _ _ -> False
      If _ stms1 stms2 -> hasYield stms1 || hasYield stms2
      IfThen _ stms -> hasYield stms
      MultiIf _ stms1 (NE.map snd -> branches) mstms2 ->
        hasYield stms1 || (or (hasYield <$> branches)) ||
          (case mstms2 of
            Nothing -> False
            Just stms -> hasYield stms)
      For _ _ _ stms _ -> hasYield stms
      ForEach _ _ stms -> hasYield stms
      While _ stms -> hasYield stms
      Loop stms -> hasYield stms
      Yield _ -> True
      _ -> False

functionStmsWithoutYield :: Gen (NonEmpty FunctionStatementWithLoc)
functionStmsWithoutYield = do
  stms <- nonEmptyGen getGen
  if hasYield stms then functionStmsWithoutYield else pure stms

functionStmsWithYield :: Gen (NonEmpty FunctionStatementWithLoc)
functionStmsWithYield = do
  stms <- nonEmptyGen getGen
  if hasYield stms then pure stms else  functionStmsWithYield