Adjust stack size based on local variables

5 files changed, 74 insertions, 48 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)
diff --git a/lib/Language/Elna/Glue.hs b/lib/Language/Elna/Glue.hs
index 9101ca5..02a9b9f 100644
--- a/lib/Language/Elna/Glue.hs
+++ b/lib/Language/Elna/Glue.hs
@@ -102,7 +102,7 @@ statement localTable (AST.CallStatement (AST.Identifier callName) arguments) = d
         $ Vector.snoc (argumentStatements <> parameterStatements)
         $ CallQuadruple callName
         $ fromIntegral
-        $ Vector.length argumentStatements
+        $ length arguments
 statement localTable (AST.CompoundStatement statements) =
     fold <$> traverse (statement localTable) statements
 statement localTable (AST.IfStatement ifCondition ifStatement elseStatement) = do
diff --git a/lib/Language/Elna/RiscV/CodeGenerator.hs b/lib/Language/Elna/RiscV/CodeGenerator.hs
index a2ff71e..15e8723 100644
--- a/lib/Language/Elna/RiscV/CodeGenerator.hs
+++ b/lib/Language/Elna/RiscV/CodeGenerator.hs
@@ -91,39 +91,45 @@ generateRiscV = flip evalState 0
     . foldlM go Vector.empty
     . HashMap.toList
   where
-    go accumulator (Identifier key, ProcedureQuadruples{ quadruples = value }) =
+    go accumulator (Identifier key, ProcedureQuadruples{ stackSize, quadruples = value }) =
         let code = Vector.cons (JumpLabel key [GlobalDirective, FunctionDirective])
-                . fold <$> mapM quadruple value
+                . fold <$> mapM (quadruple stackSize) value
          in (accumulator <>) <$> code
 
-quadruple :: RiscVQuadruple -> RiscVGenerator (Vector Statement)
-quadruple StartQuadruple = pure $ Vector.fromList
-    [ Instruction (RiscV.BaseInstruction RiscV.OpImm $ RiscV.I RiscV.SP RiscV.ADDI RiscV.SP (negate 4))
-    , Instruction (RiscV.BaseInstruction RiscV.Store $ RiscV.S 0 RiscV.SW RiscV.SP RiscV.S0)
-    , Instruction (RiscV.BaseInstruction RiscV.Store $ RiscV.S 4 RiscV.SW RiscV.SP RiscV.RA)
-    , Instruction (RiscV.BaseInstruction RiscV.OpImm $ RiscV.I RiscV.S0 RiscV.ADDI RiscV.SP 4)
-    ]
-quadruple (ParameterQuadruple operand1) =
+quadruple :: Word32 -> RiscVQuadruple -> RiscVGenerator (Vector Statement)
+quadruple stackSize StartQuadruple =
+    let totalStackSize = stackSize + 4
+     in pure $ Vector.fromList
+        [ Instruction (RiscV.BaseInstruction RiscV.OpImm $ RiscV.I RiscV.SP RiscV.ADDI RiscV.SP (negate totalStackSize))
+        , Instruction (RiscV.BaseInstruction RiscV.Store $ RiscV.S 0 RiscV.SW RiscV.SP RiscV.S0)
+        , Instruction (RiscV.BaseInstruction RiscV.Store $ RiscV.S 4 RiscV.SW RiscV.SP RiscV.RA)
+        , Instruction (RiscV.BaseInstruction RiscV.OpImm $ RiscV.I RiscV.S0 RiscV.ADDI RiscV.SP totalStackSize)
+        ]
+quadruple stackSize StopQuadruple =
+    let totalStackSize = stackSize + 4
+     in pure $ Vector.fromList
+        [ Instruction (RiscV.BaseInstruction RiscV.Load $ RiscV.I RiscV.S0 RiscV.LW RiscV.SP 0)
+        , Instruction (RiscV.BaseInstruction RiscV.Load $ RiscV.I RiscV.RA RiscV.LW RiscV.SP 4)
+        , Instruction (RiscV.BaseInstruction RiscV.OpImm $ RiscV.I RiscV.SP RiscV.ADDI RiscV.SP totalStackSize)
+        , Instruction (RiscV.BaseInstruction RiscV.Jalr $ RiscV.I RiscV.RA RiscV.JALR RiscV.Zero 0)
+        ]
+quadruple _ (ParameterQuadruple operand1) =
     let (operandRegister, statements) = loadImmediateOrRegister operand1 RiscV.A0
      in pure $ mappend statements $ Vector.fromList
         [ Instruction (RiscV.BaseInstruction RiscV.OpImm $ RiscV.I RiscV.SP RiscV.ADDI RiscV.SP (negate 4))
         , Instruction (RiscV.BaseInstruction RiscV.Store $ RiscV.S 0 RiscV.SW RiscV.SP operandRegister)
         ]
-quadruple (CallQuadruple callName numberOfArguments) = pure $ Vector.fromList
-    [ Instruction (RiscV.CallInstruction callName)
-    , Instruction (RiscV.BaseInstruction RiscV.OpImm $ RiscV.I RiscV.SP RiscV.ADDI RiscV.SP (numberOfArguments * 4))
-    ]
-quadruple StopQuadruple = pure $ Vector.fromList
-    [ Instruction (RiscV.BaseInstruction RiscV.Load $ RiscV.I RiscV.S0 RiscV.LW RiscV.SP 0)
-    , Instruction (RiscV.BaseInstruction RiscV.Load $ RiscV.I RiscV.RA RiscV.LW RiscV.SP 4)
-    , Instruction (RiscV.BaseInstruction RiscV.OpImm $ RiscV.I RiscV.SP RiscV.ADDI RiscV.SP 4)
-    , Instruction (RiscV.BaseInstruction RiscV.Jalr $ RiscV.I RiscV.RA RiscV.JALR RiscV.Zero 0)
-    ]
-quadruple (AddQuadruple operand1 operand2 store) =
+quadruple _ (CallQuadruple callName numberOfArguments) =
+    let restoreStackSize = numberOfArguments * 4
+     in pure $ Vector.fromList
+        [ Instruction (RiscV.CallInstruction callName)
+        , Instruction (RiscV.BaseInstruction RiscV.OpImm $ RiscV.I RiscV.SP RiscV.ADDI RiscV.SP restoreStackSize)
+        ]
+quadruple _ (AddQuadruple operand1 operand2 store) =
     commutativeBinary (+) RiscV.ADD (RiscV.Funct7 0b0000000) (operand1, operand2) store
-quadruple (ProductQuadruple operand1 operand2 store) =
+quadruple _ (ProductQuadruple operand1 operand2 store) =
     commutativeBinary (*) RiscV.MUL (RiscV.Funct7 0b0000001) (operand1, operand2) store
-quadruple (SubtractionQuadruple operand1 operand2 store)
+quadruple _ (SubtractionQuadruple operand1 operand2 store)
     | IntOperand immediateOperand1 <- operand1
     , IntOperand immediateOperand2 <- operand2 =
         let (storeRegister, storeStatements) = storeToStore store
@@ -158,7 +164,7 @@ quadruple (SubtractionQuadruple operand1 operand2 store)
                 $ RiscV.R storeRegister RiscV.ADD storeRegister operandRegister1
                 $ RiscV.Funct7 0b0000000
          in pure $ statements1 <> statements2 <> Vector.cons instruction storeStatements
-quadruple (NegationQuadruple operand1 store)
+quadruple _ (NegationQuadruple operand1 store)
     | IntOperand immediateOperand1 <- operand1 =
         let (storeRegister, storeStatements) = storeToStore store
          in pure $ lui (negate immediateOperand1) storeRegister <> storeStatements
@@ -170,7 +176,7 @@ quadruple (NegationQuadruple operand1 store)
                 $ RiscV.R storeRegister RiscV.SUB RiscV.Zero operandRegister1
                 $ RiscV.Funct7 0b0100000
          in pure $ statements1 <> Vector.cons instruction storeStatements
-quadruple (DivisionQuadruple operand1 operand2 store)
+quadruple _ (DivisionQuadruple operand1 operand2 store)
     | IntOperand immediateOperand1 <- operand1
     , IntOperand immediateOperand2 <- operand2 =
         if immediateOperand2 == 0
@@ -228,21 +234,21 @@ quadruple (DivisionQuadruple operand1 operand2 store)
             , JumpLabel branchLabel []
             , divisionInstruction
             ] <> storeStatements
-quadruple (LabelQuadruple (Label label)) = pure $ Vector.singleton $ JumpLabel label mempty
-quadruple (GoToQuadruple label) = pure $ Vector.singleton $ unconditionalJal label
-quadruple (EqualQuadruple operand1 operand2 goToLabel) =
+quadruple _ (LabelQuadruple (Label label)) = pure $ Vector.singleton $ JumpLabel label mempty
+quadruple _ (GoToQuadruple label) = pure $ Vector.singleton $ unconditionalJal label
+quadruple _ (EqualQuadruple operand1 operand2 goToLabel) =
     commutativeComparison (==) RiscV.BEQ (operand1, operand2) goToLabel
-quadruple (NonEqualQuadruple operand1 operand2 goToLabel) =
+quadruple _ (NonEqualQuadruple operand1 operand2 goToLabel) =
     commutativeComparison (/=) RiscV.BNE (operand1, operand2) goToLabel
-quadruple (LessQuadruple operand1 operand2 goToLabel) =
+quadruple _ (LessQuadruple operand1 operand2 goToLabel) =
     lessThan (operand1, operand2) goToLabel
-quadruple (GreaterQuadruple operand1 operand2 goToLabel) =
+quadruple _ (GreaterQuadruple operand1 operand2 goToLabel) =
     lessThan (operand2, operand1) goToLabel
-quadruple (LessOrEqualQuadruple operand1 operand2 goToLabel) =
+quadruple _ (LessOrEqualQuadruple operand1 operand2 goToLabel) =
     lessOrEqualThan (operand1, operand2) goToLabel
-quadruple (GreaterOrEqualQuadruple operand1 operand2 goToLabel) =
+quadruple _ (GreaterOrEqualQuadruple operand1 operand2 goToLabel) =
     lessOrEqualThan (operand2, operand1) goToLabel
-quadruple (AssignQuadruple operand1 store)
+quadruple _ (AssignQuadruple operand1 store)
     | IntOperand immediateOperand1 <- operand1 =
         let (storeRegister, storeStatements) = storeToStore store
          in pure $ lui immediateOperand1 storeRegister <> storeStatements
@@ -432,7 +438,7 @@ loadFromStore (RegisterStore register) = (register, mempty)
 loadFromStore (StackStore offset register) =
     let loadInstruction = Instruction
             $ RiscV.BaseInstruction RiscV.Load
-            $ RiscV.I register RiscV.LW RiscV.SP offset
+            $ RiscV.I register RiscV.LW RiscV.S0 offset
      in (register, Vector.singleton loadInstruction)
 
 storeToStore :: RiscVStore -> (RiscV.XRegister, Vector Statement)
@@ -440,5 +446,5 @@ storeToStore (RegisterStore register) = (register, mempty)
 storeToStore (StackStore offset register) =
     let storeInstruction = Instruction
             $ RiscV.BaseInstruction RiscV.Store
-            $ RiscV.S offset RiscV.SW RiscV.SP register
+            $ RiscV.S offset RiscV.SW RiscV.S0 register
      in (register, Vector.singleton storeInstruction)
diff --git a/tests/expectations/two_variables.txt b/tests/expectations/two_variables.txt
new file mode 100644
index 0000000..d763264
--- /dev/null
+++ b/tests/expectations/two_variables.txt
@@ -0,0 +1,2 @@
+58
+28
diff --git a/tests/vm/two_variables.elna b/tests/vm/two_variables.elna
new file mode 100644
index 0000000..c129efa
--- /dev/null
+++ b/tests/vm/two_variables.elna
@@ -0,0 +1,10 @@
+proc main() {
+  var i: int;
+  var j: int;
+
+  i := 58;
+  j := 28;
+
+  printi(i);
+  printi(j);
+}