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

import Data.HashMap.Strict (HashMap)
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(lift))

data Store r
    = RegisterStore r
    | StackStore Word32 r

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

newtype MachineState = MachineState
    { stackSize :: Word32
    } deriving (Eq, Show)

newtype Allocator r a = Allocator
    { runAllocator :: 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))
    -> HashMap Identifier (ProcedureQuadruples (Store r))
allocate machineConfiguration = fmap function
  where
    function :: Vector (Quadruple Variable) -> ProcedureQuadruples (Store r)
    function quadruples' =
        let (result, lastState)
                = flip runState (MachineState{ stackSize = 0 })
                $ flip runReaderT machineConfiguration
                $ runAllocator
                $ mapM quadruple quadruples'
         in ProcedureQuadruples
            { quadruples = result
            , stackSize = getField @"stackSize" lastState
            }

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

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 $ asks $ getField @"temporaryRegisters"
    pure $ RegisterStore
        $ temporaryRegisters' !! fromIntegral index
storeVariable (LocalVariable index) = do
    temporaryRegisters' <- Allocator $ asks $ getField @"temporaryRegisters"
    Allocator $ lift $ modify' $ MachineState . (+ 4) . getField @"stackSize"
    pure $ StackStore (index * 4)
        $ temporaryRegisters' !! pred (length temporaryRegisters' - fromIntegral index)