Fixed stack overflow error caused by improper cleanup of block scoped variables

rust_compiler/libs/compiler/src/v1.rs

@@ -1252,7 +1252,7 @@ impl<'a, W: std::io::Write> Compiler<'a, W> {
     fn expression_block<'v>(
         &mut self,
         mut expr: BlockExpression,
-        scope: &mut VariableScope<'v>,
+        parent_scope: &mut VariableScope<'v>,
     ) -> Result<(), Error> {
         // First, sort the expressions to ensure functions are hoisted
         expr.0.sort_by(|a, b| {
@@ -1267,10 +1267,12 @@ impl<'a, W: std::io::Write> Compiler<'a, W> {
             }
         });
 
+        let mut scope = VariableScope::scoped(parent_scope);
+
         for expr in expr.0 {
             if !self.declared_main
                 && !matches!(expr.node, Expression::Function(_))
-                && !scope.has_parent()
+                && !parent_scope.has_parent()
             {
                 self.write_output("main:")?;
                 self.declared_main = true;
@@ -1278,11 +1280,11 @@ impl<'a, W: std::io::Write> Compiler<'a, W> {
 
             match expr.node {
                 Expression::Return(ret_expr) => {
-                    self.expression_return(*ret_expr, scope)?;
+                    self.expression_return(*ret_expr, &mut scope)?;
                 }
                 _ => {
                     // Swallow errors within expressions so block can continue
-                    if let Err(e) = self.expression(expr, scope).and_then(|result| {
+                    if let Err(e) = self.expression(expr, &mut scope).and_then(|result| {
                         // If the expression was a statement that returned a temp result (e.g. `1 + 2;` line),
                         // we must free it to avoid leaking registers.
                         if let Some(comp_res) = result
@@ -1298,6 +1300,10 @@ impl<'a, W: std::io::Write> Compiler<'a, W> {
             }
         }
 
+        if scope.stack_offset() > 0 {
+            self.write_output(format!("sub sp sp {}", scope.stack_offset()))?;
+        }
+
         Ok(())
     }
 
@@ -1700,4 +1706,3 @@ impl<'a, W: std::io::Write> Compiler<'a, W> {
         Ok(())
     }
 }
-

rust_compiler/libs/compiler/src/variable_manager.rs

@@ -91,6 +91,8 @@ impl<'a> VariableScope<'a> {
     pub fn scoped(parent: &'a VariableScope<'a>) -> Self {
         Self {
             parent: Option::Some(parent),
+            temporary_vars: parent.temporary_vars.clone(),
+            persistant_vars: parent.persistant_vars.clone(),
             ..Default::default()
         }
     }
@@ -140,24 +142,32 @@ impl<'a> VariableScope<'a> {
         Ok(var_location)
     }
 
-    pub fn get_location_of(
-        &mut self,
-        var_name: impl Into<String>,
-    ) -> Result<VariableLocation, Error> {
+    pub fn get_location_of(&self, var_name: impl Into<String>) -> Result<VariableLocation, Error> {
         let var_name = var_name.into();
-        let var = self
-            .var_lookup_table
-            .get(var_name.as_str())
-            .cloned()
-            .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)
+        // 1. Check this scope
+        if let Some(var) = self.var_lookup_table.get(var_name.as_str()) {
+            if let VariableLocation::Stack(inserted_at_offset) = var {
+                // Return offset relative to CURRENT sp
+                return Ok(VariableLocation::Stack(
+                    self.stack_offset - inserted_at_offset,
+                ));
+            } else {
+                return Ok(var.clone());
+            }
         }
+
+        // 2. Recursively check parent
+        if let Some(parent) = self.parent {
+            let loc = parent.get_location_of(var_name)?;
+
+            if let VariableLocation::Stack(parent_offset) = loc {
+                return Ok(VariableLocation::Stack(parent_offset + self.stack_offset));
+            }
+            return Ok(loc);
+        }
+
+        Err(Error::UnknownVariable(var_name))
     }
 
     pub fn has_parent(&self) -> bool {