Adjust stack size based on local variables

1 files changed, 19 insertions, 11 deletions

diff --git a/lib/Language/Elna/Backend/Allocator.hs b/lib/Language/Elna/Backend/Allocator.hs
index 6c379dc..9a85605 100644
--- a/lib/Language/Elna/Backend/Allocator.hs
+++ b/lib/Language/Elna/Backend/Allocator.hs
@@ -15,8 +15,9 @@ import Language.Elna.Backend.Intermediate
     )
 import Language.Elna.Location (Identifier(..))
 import Control.Monad.Trans.Reader (ReaderT, asks, runReaderT)
-import Control.Monad.Trans.State (State, runState)
+import Control.Monad.Trans.State (State, runState, modify')
 import GHC.Records (HasField(..))
+import Control.Monad.Trans.Class (MonadTrans(lift))
 
 data Store r
     = RegisterStore r
@@ -26,8 +27,12 @@ newtype MachineConfiguration r = MachineConfiguration
     { temporaryRegisters :: [r]
     }
 
+newtype MachineState = MachineState
+    { stackSize :: Word32
+    } deriving (Eq, Show)
+
 newtype Allocator r a = Allocator
-    { runAllocator :: ReaderT (MachineConfiguration r) (State Word32) a
+    { runAllocator :: ReaderT (MachineConfiguration r) (State MachineState) a
     }
 
 instance forall r. Functor (Allocator r)
@@ -51,14 +56,16 @@ allocate
 allocate machineConfiguration = fmap function
   where
     function :: Vector (Quadruple Variable) -> ProcedureQuadruples (Store r)
-    function quadruples' = ProcedureQuadruples
-        { quadruples = fst
-            $ flip runState 0
-            $ flip runReaderT machineConfiguration
-            $ runAllocator
-            $ mapM quadruple quadruples'
-        , stackSize = 0
-        }
+    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
@@ -129,5 +136,6 @@ storeVariable (TempVariable index) = do
         $ temporaryRegisters' !! fromIntegral index
 storeVariable (LocalVariable index) = do
     temporaryRegisters' <- Allocator $ asks $ getField @"temporaryRegisters"
-    pure $ RegisterStore
+    Allocator $ lift $ modify' $ MachineState . (+ 4) . getField @"stackSize"
+    pure $ StackStore (index * 4)
         $ temporaryRegisters' !! pred (length temporaryRegisters' - fromIntegral index)