Adjust stack size based on local variables

1 files changed, 42 insertions, 36 deletions

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)