From f7b36cb81d96817d1c69ffe9025b88112de7400c Mon Sep 17 00:00:00 2001
From: Eugen Wissner <belka@caraus.de>
Date: Sun, 24 Nov 2024 13:05:11 +0100
Subject: [PATCH] Implement the while loop

---
 .ruby-version                              |  2 +-
 TODO                                       |  2 -
 lib/Language/Elna/Backend/Allocator.hs     | 60 ++++++++++++++--------
 lib/Language/Elna/Backend/Intermediate.hs  |  3 +-
 lib/Language/Elna/Frontend/AST.hs          |  6 +--
 lib/Language/Elna/Frontend/NameAnalysis.hs |  6 +--
 lib/Language/Elna/Frontend/Parser.hs       |  6 +--
 lib/Language/Elna/Frontend/TypeAnalysis.hs |  4 +-
 lib/Language/Elna/Glue.hs                  | 18 ++++---
 lib/Language/Elna/RiscV/CodeGenerator.hs   |  8 +--
 src/Main.hs                                |  9 ++--
 tests/expectations/print_after_loop.txt    |  1 +
 tests/expectations/print_in_proc.txt       |  1 +
 tests/vm/print_after_loop.elna             |  9 ++++
 tests/vm/print_in_proc.elna                |  7 +++
 15 files changed, 89 insertions(+), 53 deletions(-)
 create mode 100644 tests/expectations/print_after_loop.txt
 create mode 100644 tests/expectations/print_in_proc.txt
 create mode 100644 tests/vm/print_after_loop.elna
 create mode 100644 tests/vm/print_in_proc.elna

diff --git a/.ruby-version b/.ruby-version
index fa7adc7..9c25013 100644
--- a/.ruby-version
+++ b/.ruby-version
@@ -1 +1 @@
-3.3.5
+3.3.6
diff --git a/TODO b/TODO
index a0ab264..1fa05f5 100644
--- a/TODO
+++ b/TODO
@@ -14,8 +14,6 @@
 - Each temporary variable gets a tn register where n is the variable index. If
   there more variables the allocation will fail with out of bounds runtime
   error. Implement spill over.
-- The allocator puts temporary and local variables into the same registers,
-  causing conflicts.
 
 # Language
 
diff --git a/lib/Language/Elna/Backend/Allocator.hs b/lib/Language/Elna/Backend/Allocator.hs
index 9a85605..f0f285b 100644
--- a/lib/Language/Elna/Backend/Allocator.hs
+++ b/lib/Language/Elna/Backend/Allocator.hs
@@ -5,6 +5,7 @@ module Language.Elna.Backend.Allocator
     ) where
 
 import Data.HashMap.Strict (HashMap)
+import Data.Int (Int32)
 import Data.Word (Word32)
 import Data.Vector (Vector)
 import Language.Elna.Backend.Intermediate
@@ -17,11 +18,20 @@ import Language.Elna.Location (Identifier(..))
 import Control.Monad.Trans.Reader (ReaderT, asks, runReaderT)
 import Control.Monad.Trans.State (State, runState, modify')
 import GHC.Records (HasField(..))
-import Control.Monad.Trans.Class (MonadTrans(lift))
+import Control.Monad.Trans.Class (MonadTrans(..))
+import Control.Monad.Trans.Except (ExceptT(..), runExceptT, throwE)
+import Data.List ((!?))
 
 data Store r
     = RegisterStore r
-    | StackStore Word32 r
+    | StackStore Int32 r
+
+data AllocationError = AllocationError
+    deriving Eq
+
+instance Show AllocationError
+  where
+    show = const "Ran out of registers during register allocation"
 
 newtype MachineConfiguration r = MachineConfiguration
     { temporaryRegisters :: [r]
@@ -32,7 +42,7 @@ newtype MachineState = MachineState
     } deriving (Eq, Show)
 
 newtype Allocator r a = Allocator
-    { runAllocator :: ReaderT (MachineConfiguration r) (State MachineState) a
+    { runAllocator :: ExceptT AllocationError (ReaderT (MachineConfiguration r) (State MachineState)) a
     }
 
 instance forall r. Functor (Allocator r)
@@ -52,20 +62,22 @@ allocate
     :: forall r
     . MachineConfiguration r
     -> HashMap Identifier (Vector (Quadruple Variable))
-    -> HashMap Identifier (ProcedureQuadruples (Store r))
-allocate machineConfiguration = fmap function
+    -> Either AllocationError (HashMap Identifier (ProcedureQuadruples (Store r)))
+allocate machineConfiguration = traverse function
   where
-    function :: Vector (Quadruple Variable) -> ProcedureQuadruples (Store r)
+    run = flip runState (MachineState{ stackSize = 0 })
+        . flip runReaderT machineConfiguration
+        . runExceptT
+        . runAllocator
+        . mapM quadruple
+    function :: Vector (Quadruple Variable) -> Either AllocationError (ProcedureQuadruples (Store r))
     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
-            }
+        let (result, lastState) = run quadruples'
+         in makeResult lastState <$> result
+    makeResult MachineState{ stackSize } result = ProcedureQuadruples
+        { quadruples = result
+        , stackSize = stackSize
+        }
 
 quadruple :: Quadruple Variable -> Allocator r (Quadruple (Store r))
 quadruple = \case
@@ -131,11 +143,15 @@ operand (VariableOperand variableOperand) =
 
 storeVariable :: Variable -> Allocator r (Store r)
 storeVariable (TempVariable index) = do
-    temporaryRegisters' <- Allocator $ asks $ getField @"temporaryRegisters"
-    pure $ RegisterStore
-        $ temporaryRegisters' !! fromIntegral index
+    temporaryRegisters' <- Allocator $ lift $ asks $ getField @"temporaryRegisters"
+    maybe (Allocator $ throwE AllocationError) (pure . RegisterStore)
+        $ temporaryRegisters' !? fromIntegral index
 storeVariable (LocalVariable index) = do
-    temporaryRegisters' <- Allocator $ asks $ getField @"temporaryRegisters"
-    Allocator $ lift $ modify' $ MachineState . (+ 4) . getField @"stackSize"
-    pure $ StackStore (index * 4)
-        $ temporaryRegisters' !! pred (length temporaryRegisters' - fromIntegral index)
+    temporaryRegisters' <- Allocator $ lift $ asks $ getField @"temporaryRegisters"
+    Allocator $ lift $ lift $ modify' $ MachineState . (+ 4) . getField @"stackSize"
+    maybe (Allocator $ throwE AllocationError) (pure . StackStore (fromIntegral (succ index) * (-4)))
+        $ temporaryRegisters' !? pred (length temporaryRegisters' - fromIntegral index)
+storeVariable (ParameterVariable index) = do
+    temporaryRegisters' <- Allocator $ lift $ asks $ getField @"temporaryRegisters"
+    maybe (Allocator $ throwE AllocationError) (pure . StackStore (fromIntegral index * 4))
+        $ temporaryRegisters' !? fromIntegral index
diff --git a/lib/Language/Elna/Backend/Intermediate.hs b/lib/Language/Elna/Backend/Intermediate.hs
index 624bba8..331d044 100644
--- a/lib/Language/Elna/Backend/Intermediate.hs
+++ b/lib/Language/Elna/Backend/Intermediate.hs
@@ -19,13 +19,14 @@ instance Show Label
   where
     show (Label label) = '.' : Text.unpack label
 
-data Variable = TempVariable Word32 | LocalVariable Word32
+data Variable = TempVariable Word32 | LocalVariable Word32 | ParameterVariable Word32
     deriving Eq
 
 instance Show Variable
   where
     show (LocalVariable variable) = '@' : show variable
     show (TempVariable variable) = '$' : show variable
+    show (ParameterVariable variable) = '%' : show variable
 
 data Operand v
     = IntOperand Int32
diff --git a/lib/Language/Elna/Frontend/AST.hs b/lib/Language/Elna/Frontend/AST.hs
index e334370..df00d4b 100644
--- a/lib/Language/Elna/Frontend/AST.hs
+++ b/lib/Language/Elna/Frontend/AST.hs
@@ -71,7 +71,7 @@ data Statement
     = EmptyStatement
     | IfStatement Condition Statement (Maybe Statement)
     | AssignmentStatement VariableAccess Expression
-    -- | WhileStatement Condition Statement
+    | WhileStatement Condition Statement
     | CompoundStatement [Statement]
     | CallStatement Identifier [Expression]
     deriving Eq
@@ -86,8 +86,8 @@ instance Show Statement
         ]
     show (AssignmentStatement lhs rhs) =
         concat [show lhs, " := ", show rhs, ";"]
-    {-show (WhileStatement expression statement) =
-        concat ["while (", show expression, ") ", show statement, ";"]-}
+    show (WhileStatement expression statement) =
+        concat ["while (", show expression, ") ", show statement, ";"]
     show (CompoundStatement statements) =
         concat ["{\n", unlines (show <$> statements), " }"]
     show (CallStatement name parameters) = show name <> "("
diff --git a/lib/Language/Elna/Frontend/NameAnalysis.hs b/lib/Language/Elna/Frontend/NameAnalysis.hs
index 97a79e3..fbf5d51 100644
--- a/lib/Language/Elna/Frontend/NameAnalysis.hs
+++ b/lib/Language/Elna/Frontend/NameAnalysis.hs
@@ -175,9 +175,9 @@ statement globalTable (AST.IfStatement ifCondition ifStatement elseStatement)
 statement globalTable (AST.AssignmentStatement lvalue rvalue)
     = variableAccess globalTable lvalue
     >> expression globalTable rvalue
---statement globalTable (AST.WhileStatement whileCondition loop)
---    = condition globalTable whileCondition
---    >> statement globalTable loop
+statement globalTable (AST.WhileStatement whileCondition loop)
+    = condition globalTable whileCondition
+    >> statement globalTable loop
 
 condition :: SymbolTable -> AST.Condition -> NameAnalysis ()
 condition globalTable (AST.EqualCondition lhs rhs)
diff --git a/lib/Language/Elna/Frontend/Parser.hs b/lib/Language/Elna/Frontend/Parser.hs
index 98638f7..570c91f 100644
--- a/lib/Language/Elna/Frontend/Parser.hs
+++ b/lib/Language/Elna/Frontend/Parser.hs
@@ -188,7 +188,7 @@ statementP
     <|> ifElseP
     <|> CompoundStatement <$> blockP (many statementP)
     <|> try assignmentP
-    -- <|> try whileP
+    <|> try whileP
     <|> callP
     <?> "statement"
   where
@@ -200,9 +200,9 @@ statementP
         <$> (symbol "if" *> parensP conditionP)
         <*> statementP
         <*> optional (symbol "else" *> statementP)
-    {-whileP = WhileStatement
+    whileP = WhileStatement
         <$> (symbol "while" *> parensP conditionP)
-        <*> statementP -}
+        <*> statementP
     assignmentP = AssignmentStatement
         <$> variableAccessP
         <* symbol ":="
diff --git a/lib/Language/Elna/Frontend/TypeAnalysis.hs b/lib/Language/Elna/Frontend/TypeAnalysis.hs
index b9dcac1..c43e80e 100644
--- a/lib/Language/Elna/Frontend/TypeAnalysis.hs
+++ b/lib/Language/Elna/Frontend/TypeAnalysis.hs
@@ -101,11 +101,11 @@ statement globalTable = \case
             $ TypeAnalysis $ lift $ throwE $ InvalidConditionTypeError lhsType
         unless (rhsType == intType)
             $ TypeAnalysis $ lift $ throwE $ InvalidConditionTypeError rhsType
-    {- AST.WhileStatement whileCondition whileStatement -> do
+    AST.WhileStatement whileCondition whileStatement -> do
         conditionType <- condition globalTable whileCondition
         unless (conditionType == booleanType)
             $ TypeAnalysis $ lift $ throwE $ InvalidConditionTypeError conditionType
-        statement globalTable whileStatement -}
+        statement globalTable whileStatement
     AST.IfStatement ifCondition ifStatement elseStatement -> do
         conditionType <- condition globalTable ifCondition
         unless (conditionType == booleanType)
diff --git a/lib/Language/Elna/Glue.hs b/lib/Language/Elna/Glue.hs
index 02a9b9f..79ae461 100644
--- a/lib/Language/Elna/Glue.hs
+++ b/lib/Language/Elna/Glue.hs
@@ -71,14 +71,16 @@ declaration
     :: SymbolTable
     -> AST.Declaration
     -> Glue (Maybe (AST.Identifier, Vector (Quadruple Variable)))
-declaration globalTable (AST.ProcedureDeclaration procedureName _ variableDeclarations statements)
-    = traverse_ registerVariable variableDeclarations
+declaration globalTable (AST.ProcedureDeclaration procedureName parameters variableDeclarations statements)
+    = traverseWithIndex registerVariable variableDeclarations
+    >> traverseWithIndex registerParameter parameters
     >> nameQuadruplesTuple <$> traverse (statement globalTable) statements
   where
-    registerVariable (AST.VariableDeclaration identifier _) = do
-        currentCounter <- fmap (fromIntegral . HashMap.size)
-            $ Glue $ gets $ getField @"localMap"
-        Glue $ modify' $ modifier identifier $ LocalVariable currentCounter
+    traverseWithIndex f = traverse_ (uncurry f) . zip [0..]
+    registerParameter index (AST.Parameter identifier _ _) =
+        Glue $ modify' $ modifier identifier $ ParameterVariable index
+    registerVariable index (AST.VariableDeclaration identifier _) =
+        Glue $ modify' $ modifier identifier $ LocalVariable index
     modifier identifier currentCounter generator = generator
         { localMap = HashMap.insert identifier currentCounter
             $ getField @"localMap" generator
@@ -131,7 +133,7 @@ statement localTable (AST.AssignmentStatement variableAccess' assignee) = do
                 . AssignQuadruple rhsOperand
                 <$> lookupLocal identifier
     pure $ rhsStatements <> lhsStatements
-{- statement localTable (AST.WhileStatement whileCondition whileStatement) = do
+statement localTable (AST.WhileStatement whileCondition whileStatement) = do
     (conditionStatements, jumpConstructor) <- condition localTable whileCondition
     startLabel <- createLabel
     endLabel <- createLabel
@@ -141,7 +143,7 @@ statement localTable (AST.AssignmentStatement variableAccess' assignee) = do
         <> conditionStatements
         <> Vector.fromList [jumpConstructor startLabel, GoToQuadruple endLabel, LabelQuadruple startLabel]
         <> whileStatements
-        <> Vector.fromList [GoToQuadruple conditionLabel, LabelQuadruple endLabel] -}
+        <> Vector.fromList [GoToQuadruple conditionLabel, LabelQuadruple endLabel]
 
 createTemporary :: Glue Variable
 createTemporary = do
diff --git a/lib/Language/Elna/RiscV/CodeGenerator.hs b/lib/Language/Elna/RiscV/CodeGenerator.hs
index 15e8723..a0ad5f9 100644
--- a/lib/Language/Elna/RiscV/CodeGenerator.hs
+++ b/lib/Language/Elna/RiscV/CodeGenerator.hs
@@ -98,7 +98,7 @@ generateRiscV = flip evalState 0
 
 quadruple :: Word32 -> RiscVQuadruple -> RiscVGenerator (Vector Statement)
 quadruple stackSize StartQuadruple =
-    let totalStackSize = stackSize + 4
+    let totalStackSize = stackSize + 8
      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)
@@ -106,7 +106,7 @@ quadruple stackSize StartQuadruple =
         , Instruction (RiscV.BaseInstruction RiscV.OpImm $ RiscV.I RiscV.S0 RiscV.ADDI RiscV.SP totalStackSize)
         ]
 quadruple stackSize StopQuadruple =
-    let totalStackSize = stackSize + 4
+    let totalStackSize = stackSize + 8
      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)
@@ -438,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.S0 offset
+            $ RiscV.I register RiscV.LW RiscV.S0 (fromIntegral offset)
      in (register, Vector.singleton loadInstruction)
 
 storeToStore :: RiscVStore -> (RiscV.XRegister, Vector Statement)
@@ -446,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.S0 register
+            $ RiscV.S (fromIntegral offset) RiscV.SW RiscV.S0 register
      in (register, Vector.singleton storeInstruction)
diff --git a/src/Main.hs b/src/Main.hs
index 81e5976..0d1387f 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -25,6 +25,7 @@ import Control.Exception (IOException, catch)
 -- 3 - Parse error.
 -- 4 - Name analysis error.
 -- 5 - Type error.
+-- 6 - Register allocation error.
 
 main :: IO ()
 main = execParser commandLine >>= withCommandLine
@@ -45,9 +46,9 @@ main = execParser commandLine >>= withCommandLine
         | Just typeError <- typeAnalysis symbolTable program =
             printAndExit 5 typeError
         | otherwise =
-            let instructions = generateRiscV
-                    $ allocate riscVConfiguration
-                    $ glue symbolTable program
-             in elfObject output $ riscv32Elf instructions
+            let makeObject = elfObject output . riscv32Elf . generateRiscV
+             in either (printAndExit 6) makeObject
+                $ allocate riscVConfiguration
+                $ glue symbolTable program
     printAndExit :: Show b => forall a. Int -> b -> IO a
     printAndExit failureCode e = print e >> exitWith (ExitFailure failureCode)
diff --git a/tests/expectations/print_after_loop.txt b/tests/expectations/print_after_loop.txt
new file mode 100644
index 0000000..0cfbf08
--- /dev/null
+++ b/tests/expectations/print_after_loop.txt
@@ -0,0 +1 @@
+2
diff --git a/tests/expectations/print_in_proc.txt b/tests/expectations/print_in_proc.txt
new file mode 100644
index 0000000..8351c19
--- /dev/null
+++ b/tests/expectations/print_in_proc.txt
@@ -0,0 +1 @@
+14
diff --git a/tests/vm/print_after_loop.elna b/tests/vm/print_after_loop.elna
new file mode 100644
index 0000000..3fd9175
--- /dev/null
+++ b/tests/vm/print_after_loop.elna
@@ -0,0 +1,9 @@
+proc main() {
+  var x: int;
+
+  x := 0;
+  while (x < 2) {
+    x := x + 1;
+  }
+  printi(x);
+}
diff --git a/tests/vm/print_in_proc.elna b/tests/vm/print_in_proc.elna
new file mode 100644
index 0000000..8c4588b
--- /dev/null
+++ b/tests/vm/print_in_proc.elna
@@ -0,0 +1,7 @@
+proc print2(a: int) {
+  printi(a);
+}
+
+proc main() {
+  print2(14);
+}