diff --git a/rust_compiler/libs/compiler/src/v1.rs b/rust_compiler/libs/compiler/src/v1.rs index de7ad3a..f97e99b 100644 --- a/rust_compiler/libs/compiler/src/v1.rs +++ b/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(()) } } - diff --git a/rust_compiler/libs/compiler/src/variable_manager.rs b/rust_compiler/libs/compiler/src/variable_manager.rs index ca47f1c..c22a39f 100644 --- a/rust_compiler/libs/compiler/src/variable_manager.rs +++ b/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, - ) -> Result { + pub fn get_location_of(&self, var_name: impl Into) -> Result { 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 {