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.Word (Word32)
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 = AllocationError
    deriving Eq

instance Show AllocationError
  where
    show = const "Ran out of registers during register allocation"

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 identifier quadruples' =
        let Just (ProcedureInfo localTable _) = SymbolTable.lookup identifier globalTable
            (result, lastState) = run localTable quadruples'
         in makeResult lastState <$> result
    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 AllocationError) (pure . RegisterStore)
        $ temporaryRegisters' !? fromIntegral index
storeVariable (LocalVariable index) = do
    temporaryRegisters' <- Allocator $ lift $ asks $ getField @"temporaryRegisters"
    maybe (Allocator $ throwE AllocationError) (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 AllocationError) (pure . StackStore (fromIntegral index * 4))
        $ temporaryRegisters' !? fromIntegral index