elna/lib/Language/Elna/NameAnalysis.hs

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

import Control.Monad.Trans.Except (Except, runExcept, throwE)
import Control.Monad.Trans.Reader
    ( ReaderT(..)
    , ask
    , asks
    , runReaderT
    , withReaderT
    )
import Data.Functor ((<&>))
import qualified Language.Elna.AST as AST
import Language.Elna.Location (Identifier(..))
import Language.Elna.SymbolTable
    ( Info(..)
    , ParameterInfo(..)
    , SymbolTable
    , builtInSymbolTable
    )
import qualified Language.Elna.SymbolTable as SymbolTable
import Language.Elna.Types (Type(..))
import Control.Monad.Trans.Class (MonadTrans(..))
import Control.Monad (foldM, unless)
import qualified Data.List.NonEmpty as NonEmpty
import Data.Foldable (traverse_)
import qualified Data.Vector as Vector

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

newtype NameAnalysis a = NameAnalysis
    { runNameAnalysis :: ReaderT SymbolTable (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
    . flip runReaderT builtInSymbolTable
    . runNameAnalysis
    . program

program :: AST.Program -> NameAnalysis SymbolTable
program (AST.Program declarations)
    = NameAnalysis ask
    >>= flip (foldM declaration) declarations

declaration :: SymbolTable -> AST.Declaration -> NameAnalysis SymbolTable
declaration globalTable (AST.TypeDefinition identifier typeExpression)
    = withSymbolTable globalTable (dataType typeExpression)
    >>= flip (enter identifier) globalTable . TypeInfo
declaration globalTable (AST.ProcedureDefinition identifier parameters variables body) = do
        parametersInfo <- mapM parameter parameters
        variableInfo <- mapM variableDeclaration variables
        newTable <- either (identifierAlreadyDefinedError . NonEmpty.head) pure
            $ SymbolTable.fromList
            $ fmap parameterToVariableInfo parametersInfo
            <> variableInfo
        traverse_ (statement globalTable) body
        let procedureInfo = ProcedureInfo newTable
                $ Vector.fromList parametersInfo
        enter identifier procedureInfo globalTable

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

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

parameterToVariableInfo :: ParameterInfo -> (Identifier, Info)
parameterToVariableInfo ParameterInfo{..} =
    ( name
    , VariableInfo isReferenceParameter type'
    )

withSymbolTable :: forall a. SymbolTable -> NameAnalysis a -> NameAnalysis a
withSymbolTable symbolTable' = NameAnalysis
    . withReaderT (const symbolTable')
    . runNameAnalysis

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