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

import Data.HashMap.Strict (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, modify')
import GHC.Records (HasField(..))
import Control.Monad.Trans.Class (MonadTrans(..))
import Control.Monad.Trans.Except (ExceptT(..), runExceptT, throwE)
import Data.List ((!?))

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
    { stackSize :: Word32
    } 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
    -> HashMap Identifier (Vector (Quadruple Variable))
    -> Either AllocationError (HashMap Identifier (ProcedureQuadruples (Store r)))
allocate machineConfiguration = traverse function
  where
    run = flip runState (MachineState{ stackSize = 0 })
        . flip runReaderT machineConfiguration
        . runExceptT
        . runAllocator
        . mapM quadruple
    function :: Vector (Quadruple Variable) -> Either AllocationError (ProcedureQuadruples (Store r))
    function quadruples' =
        let (result, lastState) = run quadruples'
         in makeResult lastState <$> result
    makeResult MachineState{ stackSize } result = ProcedureQuadruples
        { quadruples = result
        , stackSize = stackSize
        }

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

operand :: Operand Variable -> Allocator r (Operand (Store r))
operand (IntOperand x) = pure $ IntOperand x
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"
    Allocator $ lift $ lift $ modify' $ MachineState . (+ 4) . getField @"stackSize"
    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