elna/lib/Language/Elna/NameAnalysis.hs

module Language.Elna.NameAnalysis
    ( nameAnalysis
    , Error(..)
    ) where

import qualified Data.List.NonEmpty as NonEmpty
import qualified Data.Vector as Vector
import qualified Language.Elna.SymbolTable as SymbolTable
import qualified Language.Elna.AST as AST
import Language.Elna.SymbolTable
    ( SymbolTable
    , Info(..)
    , ParameterInfo(..)
    )
import Control.Monad.Trans.Except (Except, runExcept, throwE)
import Data.Functor ((<&>))
import Language.Elna.Location (Identifier(..))
import Language.Elna.Types (Type(..))
import Data.Foldable (traverse_)
import Control.Monad (foldM)

data Error
    = UndefinedTypeError Identifier
    | UnexpectedTypeInfoError Info
    | IdentifierAlreadyDefinedError Identifier
    | UndefinedSymbolError Identifier
    | UnexpectedArrayByValue Identifier
    deriving (Eq, Show)

newtype NameAnalysis a = NameAnalysis
    { runNameAnalysis :: Except Error a
    }

instance Functor NameAnalysis
  where
    fmap f (NameAnalysis x) = NameAnalysis $ f <$> x

instance Applicative NameAnalysis
  where
    pure = NameAnalysis . pure
    (NameAnalysis f) <*> (NameAnalysis x) = NameAnalysis $ f <*> x

instance Monad NameAnalysis
  where
    (NameAnalysis x) >>= f = NameAnalysis $ x >>= (runNameAnalysis . f)

nameAnalysis :: AST.Program -> Either Error SymbolTable
nameAnalysis = runExcept
    . runNameAnalysis
    . program SymbolTable.builtInSymbolTable

program :: SymbolTable -> AST.Program -> NameAnalysis SymbolTable
program symbolTable (AST.Program declarations) = do
    globalTable <- foldM procedureDeclaration symbolTable declarations
    foldM declaration globalTable declarations

procedureDeclaration :: SymbolTable -> AST.Declaration -> NameAnalysis SymbolTable
procedureDeclaration globalTable (AST.ProcedureDeclaration identifier parameters _ _) = do
    parametersInfo <- mapM (parameter globalTable) parameters
    let procedureInfo = ProcedureInfo SymbolTable.empty
            $ Vector.fromList parametersInfo
    maybe (NameAnalysis $ throwE $ IdentifierAlreadyDefinedError identifier) pure
        $ SymbolTable.enter identifier procedureInfo globalTable
        
declaration :: SymbolTable -> AST.Declaration -> NameAnalysis SymbolTable
declaration globalTable (AST.ProcedureDeclaration identifier parameters variables body) = do
    variableInfo <- mapM (variableDeclaration globalTable) variables
    parameterInfo <- mapM (parameterToVariableInfo globalTable) parameters
    procedureTable <- fmap (SymbolTable.scope globalTable)
        $ either (NameAnalysis . throwE . IdentifierAlreadyDefinedError . NonEmpty.head) pure
        $ SymbolTable.fromList
        $ parameterInfo <> variableInfo
    traverse_ (statement procedureTable) body
    pure $ SymbolTable.update (updater procedureTable) identifier globalTable
  where
    updater procedureTable (ProcedureInfo _ parameters') = Just
        $ ProcedureInfo procedureTable parameters'
    updater _ _ = Nothing

parameterToVariableInfo :: SymbolTable -> AST.Parameter -> NameAnalysis (Identifier, Info)
parameterToVariableInfo symbolTable (AST.Parameter identifier typeExpression isReferenceParameter')
    = (identifier,) . VariableInfo isReferenceParameter'
    <$> dataType symbolTable typeExpression

variableDeclaration :: SymbolTable -> AST.VariableDeclaration -> NameAnalysis (Identifier, Info)
variableDeclaration globalTable (AST.VariableDeclaration identifier typeExpression)
    = (identifier,) . VariableInfo False
    <$> dataType globalTable typeExpression

parameter :: SymbolTable -> AST.Parameter -> NameAnalysis ParameterInfo
parameter environmentSymbolTable (AST.Parameter identifier typeExpression isReferenceParameter') = do
    parameterType <- dataType environmentSymbolTable typeExpression
    case parameterType of
        ArrayType _ _
            | not isReferenceParameter' -> NameAnalysis
                $ throwE $ UnexpectedArrayByValue identifier
        _ ->
            let parameterInfo = ParameterInfo
                    { name = identifier
                    , type' = parameterType
                    , isReferenceParameter = isReferenceParameter'
                    }
             in pure parameterInfo

dataType :: SymbolTable -> AST.TypeExpression -> NameAnalysis Type
dataType environmentSymbolTable (AST.NamedType baseType) = do
    case SymbolTable.lookup baseType environmentSymbolTable of
        Just baseInfo
            | TypeInfo baseType' <- baseInfo -> pure baseType'
            | otherwise -> NameAnalysis $ throwE $ UnexpectedTypeInfoError baseInfo
        _ -> NameAnalysis $ throwE $ UndefinedTypeError baseType
dataType environmentSymbolTable (AST.ArrayType arraySize baseType) =
    dataType environmentSymbolTable baseType <&> ArrayType arraySize

statement :: SymbolTable -> AST.Statement -> NameAnalysis ()
statement _ AST.EmptyStatement = pure ()
{- statement globalTable (AST.AssignmentStatement lvalue rvalue)
    = variableAccess globalTable lvalue
    >> expression globalTable rvalue
statement globalTable (AST.IfStatement ifCondition ifStatement elseStatement)
    = condition globalTable ifCondition
    >> statement globalTable ifStatement
    >> maybe (pure ()) (statement globalTable) elseStatement
statement globalTable (AST.WhileStatement whileCondition loop)
    = condition globalTable whileCondition
    >> statement globalTable loop
statement globalTable (AST.CompoundStatement statements) =
    traverse_ (statement globalTable) statements
statement globalTable (AST.CallStatement name arguments)
    = checkSymbol globalTable name
    >> traverse_ (expression globalTable) arguments

checkSymbol :: SymbolTable -> Identifier -> NameAnalysis ()
checkSymbol globalTable identifier =
    let undefinedSymbolError = NameAnalysis
            $ lift
            $ throwE
            $ UndefinedSymbolError identifier
        isDefined = SymbolTable.member identifier globalTable
     in NameAnalysis (asks (SymbolTable.member identifier))
        >>= (flip unless undefinedSymbolError . (isDefined ||))

expression :: SymbolTable -> AST.Expression -> NameAnalysis ()
expression globalTable (AST.VariableExpression variableExpression) =
    variableAccess globalTable variableExpression
expression _ (AST.LiteralExpression _) = pure ()
expression globalTable (AST.NegationExpression negation) =
    expression globalTable negation
expression globalTable (AST.SumExpression lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
expression globalTable (AST.SubtractionExpression lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
expression globalTable (AST.ProductExpression lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
expression globalTable (AST.DivisionExpression lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs

variableAccess :: SymbolTable -> AST.VariableAccess -> NameAnalysis ()
variableAccess globalTable (AST.ArrayAccess arrayExpression indexExpression)
    = variableAccess globalTable arrayExpression
    >> expression globalTable indexExpression
variableAccess globalTable (AST.VariableAccess identifier) =
    checkSymbol globalTable identifier

condition :: SymbolTable -> AST.Condition -> NameAnalysis ()
condition globalTable (AST.EqualCondition lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
condition globalTable (AST.NonEqualCondition lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
condition globalTable (AST.LessCondition lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
condition globalTable (AST.GreaterCondition lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
condition globalTable (AST.LessOrEqualCondition lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs
condition globalTable (AST.GreaterOrEqualCondition lhs rhs)
    = expression globalTable lhs
    >> expression globalTable rhs

enter :: Identifier -> Info -> SymbolTable -> NameAnalysis SymbolTable
enter identifier info table
    = maybe (identifierAlreadyDefinedError identifier) pure
    $ SymbolTable.enter identifier info table

identifierAlreadyDefinedError :: Identifier -> NameAnalysis a
identifierAlreadyDefinedError = NameAnalysis
    . lift
    . throwE
    . IdentifierAlreadyDefinedError

variableDeclaration :: AST.VariableDeclaration -> NameAnalysis (Identifier, Info)
variableDeclaration (AST.VariableDeclaration identifier typeExpression)
    = (identifier,) . VariableInfo False
    <$> dataType typeExpression
-}