{- This Source Code Form is subject to the terms of the Mozilla Public License,
   v. 2.0. If a copy of the MPL was not distributed with this file, You can
   obtain one at https://mozilla.org/MPL/2.0/. -}

module Language.Elna.Backend.Allocator
    ( MachineConfiguration(..)
    , Store(..)
    , allocate
    ) where

import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import Data.Int (Int32)
import Data.Vector (Vector)
import Language.Elna.Backend.Intermediate
    ( ProcedureQuadruples(..)
    , Operand(..)
    , Quadruple(..)
    , Variable(..)
    )
import Language.Elna.Location (Identifier(..))
import Control.Monad.Trans.Reader (ReaderT, asks, runReaderT)
import Control.Monad.Trans.State (State, runState)
import GHC.Records (HasField(..))
import Control.Monad.Trans.Class (MonadTrans(..))
import Control.Monad.Trans.Except (ExceptT(..), runExceptT, throwE)
import Data.List ((!?))
import Language.Elna.Frontend.SymbolTable (Info(..), SymbolTable)
import qualified Language.Elna.Frontend.SymbolTable as SymbolTable

data Store r
    = RegisterStore r
    | StackStore Int32 r

data AllocationError
    = OutOfRegistersError
    | UnexpectedProcedureInfoError Info
    | UndefinedSymbolError Identifier
    deriving Eq

instance Show AllocationError
  where
    show OutOfRegistersError = "Ran out of registers during register allocation"
    show (UnexpectedProcedureInfoError info) =
        "Expected to encounter a procedure, got: " <> show info
    show (UndefinedSymbolError identifier) =
        concat ["Symbol \"", show identifier, "\" is not defined"]

newtype MachineConfiguration r = MachineConfiguration
    { temporaryRegisters :: [r]
    }

newtype MachineState = MachineState
    { symbolTable :: SymbolTable
    } deriving (Eq, Show)

newtype Allocator r a = Allocator
    { runAllocator :: ExceptT AllocationError (ReaderT (MachineConfiguration r) (State MachineState)) a
    }

instance forall r. Functor (Allocator r)
  where
    fmap f = Allocator . fmap f . runAllocator

instance forall r. Applicative (Allocator r)
  where
    pure = Allocator . pure
    (Allocator x) <*> (Allocator y) = Allocator $ x <*> y

instance forall r. Monad (Allocator r)
  where
    (Allocator allocator) >>= f = Allocator $ allocator >>= (runAllocator . f)

allocate
    :: forall r
    . MachineConfiguration r
    -> SymbolTable
    -> HashMap Identifier (Vector (Quadruple Variable))
    -> Either AllocationError (HashMap Identifier (ProcedureQuadruples (Store r)))
allocate machineConfiguration globalTable = HashMap.traverseWithKey function
  where
    run localTable = flip runState (MachineState{ symbolTable = localTable })
        . flip runReaderT machineConfiguration
        . runExceptT
        . runAllocator
        . mapM quadruple
    function :: Identifier -> Vector (Quadruple Variable) -> Either AllocationError (ProcedureQuadruples (Store r))
    function procedureName quadruples' =
        case SymbolTable.lookup procedureName globalTable of
            Just (ProcedureInfo localTable _) ->
                let (result, lastState) = run localTable quadruples'
                 in makeResult lastState <$> result
            Just anotherInfo -> Left $ UnexpectedProcedureInfoError anotherInfo
            Nothing -> Left $ UndefinedSymbolError procedureName

    makeResult MachineState{ symbolTable } result = ProcedureQuadruples
        { quadruples = result
        , stackSize = fromIntegral $ SymbolTable.size symbolTable * 4
        }

quadruple :: Quadruple Variable -> Allocator r (Quadruple (Store r))
quadruple = \case
    StartQuadruple -> pure StartQuadruple
    StopQuadruple -> pure StopQuadruple
    ParameterQuadruple operand1 -> ParameterQuadruple
        <$> operand operand1
    CallQuadruple name count -> pure $ CallQuadruple name count
    AddQuadruple operand1 operand2 variable -> AddQuadruple
        <$> operand operand1
        <*> operand operand2
        <*> storeVariable variable
    SubtractionQuadruple operand1 operand2 variable -> SubtractionQuadruple
        <$> operand operand1
        <*> operand operand2
        <*> storeVariable variable
    NegationQuadruple operand1 variable -> NegationQuadruple
        <$> operand operand1
        <*> storeVariable variable
    ProductQuadruple operand1 operand2 variable -> ProductQuadruple
        <$> operand operand1
        <*> operand operand2
        <*> storeVariable variable
    DivisionQuadruple operand1 operand2 variable -> DivisionQuadruple
        <$> operand operand1
        <*> operand operand2
        <*> storeVariable variable
    LabelQuadruple label -> pure $ LabelQuadruple label
    GoToQuadruple label -> pure $ GoToQuadruple label
    EqualQuadruple operand1 operand2 goToLabel -> EqualQuadruple
        <$> operand operand1
        <*> operand operand2
        <*> pure goToLabel
    NonEqualQuadruple operand1 operand2 goToLabel -> NonEqualQuadruple
        <$> operand operand1
        <*> operand operand2
        <*> pure goToLabel
    LessQuadruple operand1 operand2 goToLabel -> LessQuadruple
        <$> operand operand1
        <*> operand operand2
        <*> pure goToLabel
    GreaterQuadruple operand1 operand2 goToLabel -> do
        operand1' <- operand operand1
        operand2' <- operand operand2
        pure $ GreaterQuadruple operand1' operand2' goToLabel
    LessOrEqualQuadruple operand1 operand2 goToLabel -> LessOrEqualQuadruple
        <$> operand operand1
        <*> operand operand2
        <*> pure goToLabel
    GreaterOrEqualQuadruple operand1 operand2 goToLabel -> GreaterOrEqualQuadruple
        <$> operand operand1
        <*> operand operand2
        <*> pure goToLabel
    AssignQuadruple operand1 variable -> AssignQuadruple 
        <$> operand operand1
        <*> storeVariable variable
    ArrayAssignQuadruple operand1 operand2 variable -> ArrayAssignQuadruple 
        <$> operand operand1
        <*> operand operand2
        <*> storeVariable variable
    ArrayQuadruple variable1 operand1 variable2 -> ArrayQuadruple 
        <$> storeVariable variable1
        <*> operand operand1
        <*> storeVariable variable2

operand :: Operand Variable -> Allocator r (Operand (Store r))
operand (IntOperand literalOperand) = pure $ IntOperand literalOperand
operand (VariableOperand variableOperand) =
    VariableOperand <$> storeVariable variableOperand

storeVariable :: Variable -> Allocator r (Store r)
storeVariable (TempVariable index) = do
    temporaryRegisters' <- Allocator $ lift $ asks $ getField @"temporaryRegisters"
    maybe (Allocator $ throwE OutOfRegistersError) (pure . RegisterStore)
        $ temporaryRegisters' !? fromIntegral index
storeVariable (LocalVariable index) = do
    temporaryRegisters' <- Allocator $ lift $ asks $ getField @"temporaryRegisters"
    maybe (Allocator $ throwE OutOfRegistersError) (pure . StackStore (fromIntegral (succ index) * (-4)))
        $ temporaryRegisters' !? pred (length temporaryRegisters' - fromIntegral index)
storeVariable (ParameterVariable index) = do
    temporaryRegisters' <- Allocator $ lift $ asks $ getField @"temporaryRegisters"
    maybe (Allocator $ throwE OutOfRegistersError) (pure . StackStore (fromIntegral index * 4))
        $ temporaryRegisters' !? fromIntegral index