Split the allocator and add state to it

lib/Language/Elna/Backend/Allocator.hs

@@ -14,6 +14,9 @@ import Language.Elna.Backend.Intermediate
     , 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(..))
 
 data Store r
     = RegisterStore r
@@ -23,61 +26,108 @@ newtype MachineConfiguration r = MachineConfiguration
     { temporaryRegisters :: [r]
     }
 
+newtype Allocator r a = Allocator
+    { runAllocator :: ReaderT (MachineConfiguration r) (State Word32) 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
+allocate machineConfiguration = fmap function
   where
     function :: Vector (Quadruple Variable) -> ProcedureQuadruples (Store r)
     function quadruples' = ProcedureQuadruples
-        { quadruples = quadruple <$> quadruples'
+        { quadruples = fst
+            $ flip runState 0
+            $ flip runReaderT machineConfiguration
+            $ runAllocator
+            $ mapM quadruple quadruples'
         , stackSize = 0
         }
-    quadruple :: Quadruple Variable -> Quadruple (Store r)
+
+quadruple :: Quadruple Variable -> Allocator r (Quadruple (Store r))
 quadruple = \case
-        StartQuadruple -> StartQuadruple
-        StopQuadruple -> StopQuadruple
-        ParameterQuadruple operand1 ->
-            ParameterQuadruple (operand operand1)
-        CallQuadruple name count -> 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 -> LabelQuadruple label
-        GoToQuadruple label -> GoToQuadruple label
-        EqualQuadruple operand1 operand2 goToLabel ->
-            EqualQuadruple (operand operand1) (operand operand2) goToLabel
-        NonEqualQuadruple operand1 operand2 goToLabel ->
-            NonEqualQuadruple (operand operand1) (operand operand2) goToLabel
-        LessQuadruple operand1 operand2 goToLabel ->
-            LessQuadruple (operand operand1) (operand operand2) goToLabel
-        GreaterQuadruple operand1 operand2 goToLabel ->
-            GreaterQuadruple (operand operand1) (operand operand2) goToLabel
-        LessOrEqualQuadruple operand1 operand2 goToLabel ->
-            LessOrEqualQuadruple (operand operand1) (operand operand2) goToLabel
-        GreaterOrEqualQuadruple operand1 operand2 goToLabel ->
-            GreaterOrEqualQuadruple (operand operand1) (operand operand2) goToLabel
-        AssignQuadruple operand1 variable ->
-            AssignQuadruple (operand operand1) $ storeVariable variable
-    operand :: Operand Variable -> Operand (Store r)
-    operand (IntOperand x) = IntOperand x
+    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 (TempVariable index) = RegisterStore
-        $ temporaryRegisters !! fromIntegral index
-    storeVariable (LocalVariable index) = RegisterStore
-        $ temporaryRegisters !! pred (length temporaryRegisters - fromIntegral index)
+    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"
+    pure $ RegisterStore
+        $ temporaryRegisters' !! pred (length temporaryRegisters' - fromIntegral index)