Calculate stack offset from array index

4 files changed, 42 insertions, 33 deletions

diff --git a/lib/Language/Elna/Backend/Allocator.hs b/lib/Language/Elna/Backend/Allocator.hs
index df50d6e..15e8b6f 100644
--- a/lib/Language/Elna/Backend/Allocator.hs
+++ b/lib/Language/Elna/Backend/Allocator.hs
@@ -7,7 +7,6 @@ module Language.Elna.Backend.Allocator
 import Data.HashMap.Strict (HashMap)
 import qualified Data.HashMap.Strict as HashMap
 import Data.Int (Int32)
-import Data.Word (Word32)
 import Data.Vector (Vector)
 import Language.Elna.Backend.Intermediate
     ( ProcedureQuadruples(..)
diff --git a/lib/Language/Elna/RiscV/CodeGenerator.hs b/lib/Language/Elna/RiscV/CodeGenerator.hs
index a1dcdbe..fffbd39 100644
--- a/lib/Language/Elna/RiscV/CodeGenerator.hs
+++ b/lib/Language/Elna/RiscV/CodeGenerator.hs
@@ -276,25 +276,14 @@ quadruple _ (ArrayAssignQuadruple assigneeOperand indexOperand store)
     storeWithOffset (StackStore offset register) (IntOperand indexOffset) =
         let storeInstruction = Instruction
                 $ RiscV.BaseInstruction RiscV.Store
-                $ RiscV.S (fromIntegral $ offset + indexOffset) RiscV.SW RiscV.S0 register
-        in (register, Vector.singleton storeInstruction)
+                $ RiscV.S (fromIntegral $ offset + indexOffset * 4) RiscV.SW RiscV.S0 register
+         in (register, Vector.singleton storeInstruction)
     storeWithOffset (StackStore offset register) (VariableOperand indexOffset) =
-        let baseRegisterInstruction = Instruction
-                $ RiscV.BaseInstruction RiscV.OpImm
-                $ RiscV.I immediateRegister RiscV.ADDI RiscV.S0 0
-            (indexRegister, indexStatements) = loadFromStore indexOffset
-            registerWithOffset = Instruction
-                $ RiscV.BaseInstruction RiscV.OpImm
-                $ RiscV.I immediateRegister RiscV.ADDI indexRegister 0
-            storeInstruction = Instruction
+        let storeInstruction = Instruction
                 $ RiscV.BaseInstruction RiscV.Store
                 $ RiscV.S (fromIntegral offset) RiscV.SW immediateRegister register
-            statements = Vector.fromList
-                [ baseRegisterInstruction
-                , registerWithOffset
-                , storeInstruction
-                ]
-        in (register, indexStatements <> statements)
+            statements = calculateIndexOffset indexOffset
+         in (register, Vector.snoc statements storeInstruction)
 quadruple _ (ArrayQuadruple assigneeVariable indexOperand store) =
     let (operandRegister1, statements1) = loadWithOffset assigneeVariable indexOperand
         (storeRegister, storeStatements) = storeToStore store
@@ -308,25 +297,33 @@ quadruple _ (ArrayQuadruple assigneeVariable indexOperand store) =
     loadWithOffset (StackStore offset register) (IntOperand indexOffset) =
         let loadInstruction = Instruction
                 $ RiscV.BaseInstruction RiscV.Load
-                $ RiscV.I register RiscV.LW RiscV.S0 (fromIntegral $ offset + indexOffset)
-        in (register, Vector.singleton loadInstruction)
+                $ RiscV.I register RiscV.LW RiscV.S0 (fromIntegral $ offset + indexOffset * 4)
+         in (register, Vector.singleton loadInstruction)
     loadWithOffset (StackStore offset register) (VariableOperand indexOffset) =
-        let baseRegisterInstruction = Instruction
-                $ RiscV.BaseInstruction RiscV.OpImm
-                $ RiscV.I immediateRegister RiscV.ADDI RiscV.S0 0
-            (indexRegister, indexStatements) = loadFromStore indexOffset
-            registerWithOffset = Instruction
-                $ RiscV.BaseInstruction RiscV.OpImm
-                $ RiscV.I immediateRegister RiscV.ADDI indexRegister 0
-            loadInstruction = Instruction
+        let loadInstruction = Instruction
                 $ RiscV.BaseInstruction RiscV.Load
                 $ RiscV.I register RiscV.SW immediateRegister (fromIntegral offset)
-            statements = Vector.fromList
-                [ baseRegisterInstruction
-                , registerWithOffset
-                , loadInstruction
-                ]
-        in (register, indexStatements <> statements)
+            statements = calculateIndexOffset indexOffset
+         in (register, Vector.snoc statements loadInstruction)
+
+calculateIndexOffset :: RiscVStore -> Vector Statement
+calculateIndexOffset indexOffset =
+    let (indexRegister, indexStatements) = loadFromStore indexOffset
+        baseRegisterInstruction = Instruction
+            $ RiscV.BaseInstruction RiscV.OpImm
+            $ RiscV.I immediateRegister RiscV.ADDI RiscV.Zero 4
+        indexRelativeOffset = Instruction
+            $ RiscV.BaseInstruction RiscV.Op
+            $ RiscV.R immediateRegister RiscV.MUL immediateRegister indexRegister (RiscV.Funct7 0b0000001)
+        registerWithOffset = Instruction
+            $ RiscV.BaseInstruction RiscV.Op
+            $ RiscV.R immediateRegister RiscV.ADD immediateRegister RiscV.S0 (RiscV.Funct7 0b0000000)
+        statements = Vector.fromList
+            [ baseRegisterInstruction
+            , indexRelativeOffset
+            , registerWithOffset
+            ]
+     in indexStatements <> statements
 
 unconditionalJal :: Label -> Statement
 unconditionalJal (Label goToLabel) = Instruction
diff --git a/tests/expectations/array_with_variable_index.txt b/tests/expectations/array_with_variable_index.txt
new file mode 100644
index 0000000..b3172d1
--- /dev/null
+++ b/tests/expectations/array_with_variable_index.txt
@@ -0,0 +1,2 @@
+5
+7
diff --git a/tests/vm/array_with_variable_index.elna b/tests/vm/array_with_variable_index.elna
new file mode 100644
index 0000000..c3449b3
--- /dev/null
+++ b/tests/vm/array_with_variable_index.elna
@@ -0,0 +1,11 @@
+proc main() {
+  var a: array[2] of int;
+  var i: int;
+
+  i := 1;
+  a[0] := 5;
+  a[i] := 7;
+
+  printi(a[0]);
+  printi(a[i]);
+}