Extraction of function arguments from stack

libs/compiler/src/test/mod.rs

@@ -59,9 +59,11 @@ fn test_function_declaration_with_spillover_params() -> anyhow::Result<()> {
             pop r14 #arg3
             pop r15 #arg2
             push ra
-            pop ra
-            sub sp 1
+            sub r0 sp 1
+            get ra db r0
+            sub sp sp 2
             j ra
+            main:
         "}
     );
 
@@ -84,8 +86,11 @@ fn test_function_declaration_with_register_params() -> anyhow::Result<()> {
             pop r8 #arg2
             pop r9 #arg1
             push ra
-            pop ra
+            sub r0 sp 1
+            get ra db r0
+            sub sp sp 1
             j ra
+            main:
         "}
     );
 

libs/compiler/src/v2.rs

@@ -85,9 +85,7 @@ impl<'a, W: std::io::Write> Compiler<'a, W> {
     ) -> Result<(), Error> {
         match expr {
             Expression::Function(expr_func) => self.expression_function(expr_func, scope)?,
-            Expression::Block(expr_block) => {
-                self.expression_block(expr_block, &mut VariableScope::scoped(&scope))?
-            }
+            Expression::Block(expr_block) => self.expression_block(expr_block, scope)?,
             _ => todo!(),
         };
 
@@ -111,7 +109,10 @@ impl<'a, W: std::io::Write> Compiler<'a, W> {
         });
 
         for expr in expr.0 {
-            if !self.declared_main && !matches!(expr, Expression::Function(_)) {
+            if !self.declared_main
+                && !matches!(expr, Expression::Function(_))
+                && !scope.has_parent()
+            {
                 self.write_output("main:")?;
                 self.declared_main = true;
             }
@@ -192,11 +193,21 @@ impl<'a, W: std::io::Write> Compiler<'a, W> {
         }
 
         self.write_output("push ra")?;
+        block_scope.add_variable(format!("{name}_ra"), LocationRequest::Stack)?;
+
         self.expression_block(body, &mut block_scope)?;
-        self.write_output("pop ra")?;
+        // Get the saved return address and save it back into `ra`
+        let VariableLocation::Stack(ra_stack_offset) =
+            block_scope.get_location_of(format!("{name}_ra"))?
+        else {
+            panic!("This shouldn't happen");
+        };
+
+        self.write_output(format!("sub r0 sp {ra_stack_offset}"))?;
+        self.write_output("get ra db r0")?;
 
         if block_scope.stack_offset() > 0 {
-            self.write_output(format!("sub sp {}", block_scope.stack_offset()))?;
+            self.write_output(format!("sub sp sp {}", block_scope.stack_offset()))?;
         }
 
         self.write_output("j ra")?;

libs/compiler/src/variable_manager.rs

@@ -1,10 +1,11 @@
-// r0  - r7  : temporary variables
-// r8 - r15 : persistant variables
+// r0        : Return Value / Temp Stack Pointer
+// r1 - r7   : Temporary Variables
+// r8 - r15  : Persistant Variables
 
 use quick_error::quick_error;
 use std::collections::{HashMap, VecDeque};
 
-const TEMP: [u8; 8] = [0, 1, 2, 3, 4, 5, 6, 7];
+const TEMP: [u8; 7] = [1, 2, 3, 4, 5, 6, 7];
 const PERSIST: [u8; 8] = [8, 9, 10, 11, 12, 13, 14, 15];
 
 quick_error! {
@@ -137,10 +138,23 @@ impl<'a> VariableScope<'a> {
         var_name: impl Into<String>,
     ) -> Result<VariableLocation, Error> {
         let var_name = var_name.into();
-        self.var_lookup_table
+        let var = self
+            .var_lookup_table
             .get(var_name.as_str())
             .map(|v| v.clone())
-            .ok_or(Error::UnknownVariable(var_name))
+            .ok_or(Error::UnknownVariable(var_name))?;
+
+        if let VariableLocation::Stack(inserted_at_offset) = var {
+            Ok(VariableLocation::Stack(
+                self.stack_offset - inserted_at_offset,
+            ))
+        } else {
+            Ok(var)
+        }
+    }
+
+    pub fn has_parent(&self) -> bool {
+        self.parent.is_some()
     }
 
     pub fn free_temp(&mut self, var_name: impl Into<String>) -> Result<(), Error> {