Fix function invocation stack underflow

Changelog.md

@@ -1,5 +1,10 @@
 # Changelog
 
+[0.2.3]
+
+- Fixed stack underflow with function invocations
+  - They are still "heavy", but they should work as expected now
+
 [0.2.2]
 
 - Fixed some formatting issues when converting Markdown to Text Mesh Pro for

ModData/About/About.xml

@@ -2,7 +2,7 @@
 <ModMetadata xmlns:xsd="http://www.w3.org/2001/XMLSchema" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance">
   <Name>Slang</Name>
   <Author>JoeDiertay</Author>
-  <Version>0.2.2</Version>
+  <Version>0.2.3</Version>
   <Description>
 [h1]Slang: High-Level Programming for Stationeers[/h1]
 

csharp_mod/stationeersSlang.csproj

@@ -5,7 +5,7 @@
     <Nullable>enable</Nullable>
     <AssemblyName>StationeersSlang</AssemblyName>
     <Description>Slang Compiler Bridge</Description>
-    <Version>0.2.2</Version>
+    <Version>0.2.3</Version>
     <AllowUnsafeBlocks>true</AllowUnsafeBlocks>
     <LangVersion>latest</LangVersion>
   </PropertyGroup>

rust_compiler/Cargo.lock

@@ -909,7 +909,7 @@ checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e"
 
 [[package]]
 name = "slang"
-version = "0.2.2"
+version = "0.2.3"
 dependencies = [
  "anyhow",
  "clap",

rust_compiler/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "slang"
-version = "0.2.2"
+version = "0.2.3"
 edition = "2021"
 
 [workspace]

rust_compiler/libs/compiler/src/test/binary_expression.rs

@@ -52,6 +52,8 @@ fn nested_binary_expressions() -> Result<()> {
             add r1 r10 r9
             mul r2 r1 r8
             move r15 r2
+            j L1
+            L1:
             sub r0 sp 1
             get ra db r0
             sub sp sp 1

rust_compiler/libs/compiler/src/test/declaration_function_invocation.rs

@@ -17,6 +17,7 @@ fn no_arguments() -> anyhow::Result<()> {
         j main
         doSomething:
         push ra
+        L1:
         sub r0 sp 1
         get ra db r0
         sub sp sp 1
@@ -34,14 +35,17 @@ fn no_arguments() -> anyhow::Result<()> {
 
 #[test]
 fn let_var_args() -> anyhow::Result<()> {
-    // !IMPORTANT this needs to be stabilized as it currently incorrectly calculates sp offset at
-    // both ends of the cleanup lifecycle
     let compiled = compile! {
         debug
         "
-        fn doSomething(arg1) {};
+        fn mul2(arg1) {
-        let arg1 = 123;
+            return arg1 * 2;
-        let i = doSomething(arg1);
+        };
+        loop {
+            let arg1 = 123;
+            let i = mul2(arg1);
+            i = i ** 2;
+        }
         "
     };
 
@@ -50,23 +54,31 @@ fn let_var_args() -> anyhow::Result<()> {
         indoc! {
             "
             j main
-            doSomething:
+            mul2:
             pop r8 #arg1
             push ra
+            mul r1 r8 2
+            move r15 r1
+            j L1
+            L1:
             sub r0 sp 1
             get ra db r0
             sub sp sp 1
             j ra
             main:
+            L2:
             move r8 123 #arg1
             push r8
             push r8
-            jal doSomething
+            jal mul2
             sub r0 sp 1
             get r8 db r0
             sub sp sp 1
             move r9 r15 #i
-            sub sp sp 1
+            pow r1 r9 2
+            move r9 r1 #i
+            j L2
+            L3:
             "
         }
     );
@@ -97,7 +109,9 @@ fn inline_literal_args() -> anyhow::Result<()> {
     let compiled = compile! {
         debug
         "
-        fn doSomething(arg1, arg2) {};
+        fn doSomething(arg1, arg2) {
+            return 5;
+        };
         let thisVariableShouldStayInPlace = 123;
         let returnedValue = doSomething(12, 34);
         "
@@ -112,6 +126,9 @@ fn inline_literal_args() -> anyhow::Result<()> {
             pop r8 #arg2
             pop r9 #arg1
             push ra
+            move r15 5 #returnValue
+            j L1
+            L1:
             sub r0 sp 1
             get ra db r0
             sub sp sp 1
@@ -126,7 +143,6 @@ fn inline_literal_args() -> anyhow::Result<()> {
             get r8 db r0
             sub sp sp 1
             move r9 r15 #returnedValue
-            sub sp sp 1
             "
         }
     );
@@ -154,6 +170,7 @@ fn mixed_args() -> anyhow::Result<()> {
             pop r8 #arg2
             pop r9 #arg1
             push ra
+            L1:
             sub r0 sp 1
             get ra db r0
             sub sp sp 1
@@ -168,7 +185,6 @@ fn mixed_args() -> anyhow::Result<()> {
             get r8 db r0
             sub sp sp 1
             move r9 r15 #returnValue
-            sub sp sp 1
             "
         }
     );
@@ -198,6 +214,8 @@ fn with_return_statement() -> anyhow::Result<()> {
             pop r8 #arg1
             push ra
             move r15 456 #returnValue
+            j L1
+            L1:
             sub r0 sp 1
             get ra db r0
             sub sp sp 1
@@ -233,6 +251,7 @@ fn with_negative_return_literal() -> anyhow::Result<()> {
             doSomething:
             push ra
             move r15 -1 #returnValue
+            L1:
             sub r0 sp 1
             get ra db r0
             sub sp sp 1

rust_compiler/libs/compiler/src/test/declaration_literal.rs

@@ -120,7 +120,7 @@ fn test_boolean_return() -> anyhow::Result<()> {
         fn getTrue() {
             return true;
         };
-        
+
         let val = getTrue();
         "
     };
@@ -133,6 +133,8 @@ fn test_boolean_return() -> anyhow::Result<()> {
             getTrue:
             push ra
             move r15 1 #returnValue
+            j L1
+            L1:
             sub r0 sp 1
             get ra db r0
             sub sp sp 1

rust_compiler/libs/compiler/src/test/function_declaration.rs

@@ -21,6 +21,7 @@ fn test_function_declaration_with_spillover_params() -> anyhow::Result<()> {
             pop r13 #arg4
             pop r14 #arg3
             push ra
+            L1:
             sub r0 sp 1
             get ra db r0
             sub sp sp 3
@@ -31,6 +32,48 @@ fn test_function_declaration_with_spillover_params() -> anyhow::Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_early_return() -> anyhow::Result<()> {
+    let compiled = compile!(debug r#"
+        // This is a test function declaration with no body
+        fn doSomething() {
+            if (1 == 1) {
+                return;
+            }
+            let i = 1 + 2;
+            return;
+        };
+        doSomething();
+    "#);
+
+    assert_eq!(
+        compiled,
+        indoc! {
+            "
+            j main
+            doSomething:
+            push ra
+            seq r1 1 1
+            beq r1 0 L2
+            j L1
+            L2:
+            move r8 3 #i
+            j L1
+            L1:
+            sub r0 sp 1
+            get ra db r0
+            sub sp sp 1
+            j ra
+            main:
+            jal doSomething
+            move r1 r15 #__binary_temp_2
+            "
+        }
+    );
+
+    Ok(())
+}
+
 #[test]
 fn test_function_declaration_with_register_params() -> anyhow::Result<()> {
     let compiled = compile!(debug r#"
@@ -47,6 +90,7 @@ fn test_function_declaration_with_register_params() -> anyhow::Result<()> {
             pop r8 #arg2
             pop r9 #arg1
             push ra
+            L1:
             sub r0 sp 1
             get ra db r0
             sub sp sp 1

rust_compiler/libs/compiler/src/v1.rs

@@ -165,6 +165,7 @@ pub struct Compiler<'a, 'w, W: std::io::Write> {
     temp_counter: usize,
     label_counter: usize,
     loop_stack: Vec<(Cow<'a, str>, Cow<'a, str>)>, // Stores (start_label, end_label)
+    current_return_label: Option<Cow<'a, str>>,
     pub errors: Vec<Error<'a>>,
 }
 
@@ -186,6 +187,7 @@ impl<'a, 'w, W: std::io::Write> Compiler<'a, 'w, W> {
             temp_counter: 0,
             label_counter: 0,
             loop_stack: Vec::new(),
+            current_return_label: None,
             errors: Vec::new(),
         }
     }
@@ -900,7 +902,7 @@ impl<'a, 'w, W: std::io::Write> Compiler<'a, 'w, W> {
     fn expression_function_invocation(
         &mut self,
         invoke_expr: Spanned<InvocationExpression<'a>>,
-        stack: &mut VariableScope<'a, '_>,
+        parent_scope: &mut VariableScope<'a, '_>,
     ) -> Result<(), Error<'a>> {
         let InvocationExpression { name, arguments } = invoke_expr.node;
 
@@ -926,9 +928,10 @@ impl<'a, 'w, W: std::io::Write> Compiler<'a, 'w, W> {
             // Best to skip generation of this call to prevent bad IC10
             return Ok(());
         }
+        let mut stack = VariableScope::scoped(parent_scope);
 
         // backup all used registers to the stack
-        let active_registers = stack.registers().cloned().collect::<Vec<_>>();
+        let active_registers = stack.registers();
         for register in &active_registers {
             stack.add_variable(
                 Cow::from(format!("temp_{register}")),
@@ -991,7 +994,7 @@ impl<'a, 'w, W: std::io::Write> Compiler<'a, 'w, W> {
                 }
                 Expression::Binary(bin_expr) => {
                     // Compile the binary expression to a temp register
-                    let result = self.expression_binary(bin_expr, stack)?;
+                    let result = self.expression_binary(bin_expr, &mut stack)?;
                     let reg_str = self.resolve_register(&result.location)?;
                     self.write_output(format!("push {reg_str}"))?;
                     if let Some(name) = result.temp_name {
@@ -1000,7 +1003,7 @@ impl<'a, 'w, W: std::io::Write> Compiler<'a, 'w, W> {
                 }
                 Expression::Logical(log_expr) => {
                     // Compile the logical expression to a temp register
-                    let result = self.expression_logical(log_expr, stack)?;
+                    let result = self.expression_logical(log_expr, &mut stack)?;
                     let reg_str = self.resolve_register(&result.location)?;
                     self.write_output(format!("push {reg_str}"))?;
                     if let Some(name) = result.temp_name {
@@ -1019,7 +1022,7 @@ impl<'a, 'w, W: std::io::Write> Compiler<'a, 'w, W> {
                                 end_line: 0,
                             }, // Dummy span
                         },
-                        stack,
+                        &mut stack,
                     )?;
 
                     if let Some(result) = result_opt {
@@ -1592,7 +1595,7 @@ impl<'a, 'w, W: std::io::Write> Compiler<'a, 'w, W> {
 
             match expr.node {
                 Expression::Return(ret_expr) => {
-                    self.expression_return(*ret_expr, &mut scope)?;
+                    self.expression_return(ret_expr, &mut scope)?;
                 }
                 _ => {
                     // Swallow errors within expressions so block can continue
@@ -1622,133 +1625,149 @@ impl<'a, 'w, W: std::io::Write> Compiler<'a, 'w, W> {
     /// Takes the result of the expression and stores it in VariableScope::RETURN_REGISTER
     fn expression_return(
         &mut self,
-        expr: Spanned<Expression<'a>>,
+        expr: Option<Box<Spanned<Expression<'a>>>>,
         scope: &mut VariableScope<'a, '_>,
     ) -> Result<VariableLocation<'a>, Error<'a>> {
-        if let Expression::Negation(neg_expr) = &expr.node
+        if let Some(expr) = expr {
-            && let Expression::Literal(spanned_lit) = &neg_expr.node
+            if let Expression::Negation(neg_expr) = &expr.node
-            && let Literal::Number(neg_num) = &spanned_lit.node
+                && let Expression::Literal(spanned_lit) = &neg_expr.node
-        {
+                && let Literal::Number(neg_num) = &spanned_lit.node
-            let loc = VariableLocation::Persistant(VariableScope::RETURN_REGISTER);
+            {
-            self.emit_variable_assignment(
+                let loc = VariableLocation::Persistant(VariableScope::RETURN_REGISTER);
-                Cow::from("returnValue"),
+                self.emit_variable_assignment(
-                &loc,
+                    Cow::from("returnValue"),
-                Cow::from(format!("-{neg_num}")),
+                    &loc,
-            )?;
+                    Cow::from(format!("-{neg_num}")),
-            return Ok(loc);
-        };
-
-        match expr.node {
-            Expression::Variable(var_name) => {
-                match scope.get_location_of(&var_name.node, Some(var_name.span)) {
-                    Ok(loc) => match loc {
-                        VariableLocation::Temporary(reg) | VariableLocation::Persistant(reg) => {
-                            self.write_output(format!(
-                                "move r{} r{reg} {}",
-                                VariableScope::RETURN_REGISTER,
-                                debug!(self, "#returnValue")
-                            ))?;
-                        }
-                        VariableLocation::Constant(lit) => {
-                            let str = extract_literal(lit, false)?;
-                            self.write_output(format!(
-                                "move r{} {str} {}",
-                                VariableScope::RETURN_REGISTER,
-                                debug!(self, "#returnValue")
-                            ))?
-                        }
-                        VariableLocation::Stack(offset) => {
-                            self.write_output(format!(
-                                "sub r{} sp {offset}",
-                                VariableScope::TEMP_STACK_REGISTER
-                            ))?;
-                            self.write_output(format!(
-                                "get r{} db r{}",
-                                VariableScope::RETURN_REGISTER,
-                                VariableScope::TEMP_STACK_REGISTER
-                            ))?;
-                        }
-                        VariableLocation::Device(_) => {
-                            return Err(Error::Unknown(
-                                "You can not return a device from a function.".into(),
-                                Some(var_name.span),
-                            ));
-                        }
-                    },
-                    Err(_) => {
-                        self.errors.push(Error::UnknownIdentifier(
-                            var_name.node.clone(),
-                            var_name.span,
-                        ));
-                        // Proceed with dummy
-                    }
-                }
-            }
-            Expression::Literal(spanned_lit) => match spanned_lit.node {
-                Literal::Number(num) => {
-                    self.emit_variable_assignment(
-                        Cow::from("returnValue"),
-                        &VariableLocation::Persistant(VariableScope::RETURN_REGISTER),
-                        Cow::from(num.to_string()),
-                    )?;
-                }
-                Literal::Boolean(b) => {
-                    let val = if b { "1" } else { "0" };
-                    self.emit_variable_assignment(
-                        Cow::from("returnValue"),
-                        &VariableLocation::Persistant(VariableScope::RETURN_REGISTER),
-                        Cow::from(val.to_string()),
-                    )?;
-                }
-                _ => {}
-            },
-            Expression::Binary(bin_expr) => {
-                let result = self.expression_binary(bin_expr, scope)?;
-                let result_reg = self.resolve_register(&result.location)?;
-                self.write_output(format!(
-                    "move r{} {}",
-                    VariableScope::RETURN_REGISTER,
-                    result_reg
-                ))?;
-                if let Some(name) = result.temp_name {
-                    scope.free_temp(name, None)?;
-                }
-            }
-            Expression::Logical(log_expr) => {
-                let result = self.expression_logical(log_expr, scope)?;
-                let result_reg = self.resolve_register(&result.location)?;
-                self.write_output(format!(
-                    "move r{} {}",
-                    VariableScope::RETURN_REGISTER,
-                    result_reg
-                ))?;
-                if let Some(name) = result.temp_name {
-                    scope.free_temp(name, None)?;
-                }
-            }
-            Expression::MemberAccess(access) => {
-                // Return result of member access
-                let res_opt = self.expression(
-                    Spanned {
-                        node: Expression::MemberAccess(access),
-                        span: expr.span,
-                    },
-                    scope,
                 )?;
-                if let Some(res) = res_opt {
+                return Ok(loc);
-                    let reg = self.resolve_register(&res.location)?;
+            };
-                    self.write_output(format!("move r{} {}", VariableScope::RETURN_REGISTER, reg))?;
+
-                    if let Some(temp) = res.temp_name {
+            match expr.node {
-                        scope.free_temp(temp, None)?;
+                Expression::Variable(var_name) => {
+                    match scope.get_location_of(&var_name.node, Some(var_name.span)) {
+                        Ok(loc) => match loc {
+                            VariableLocation::Temporary(reg)
+                            | VariableLocation::Persistant(reg) => {
+                                self.write_output(format!(
+                                    "move r{} r{reg} {}",
+                                    VariableScope::RETURN_REGISTER,
+                                    debug!(self, "#returnValue")
+                                ))?;
+                            }
+                            VariableLocation::Constant(lit) => {
+                                let str = extract_literal(lit, false)?;
+                                self.write_output(format!(
+                                    "move r{} {str} {}",
+                                    VariableScope::RETURN_REGISTER,
+                                    debug!(self, "#returnValue")
+                                ))?
+                            }
+                            VariableLocation::Stack(offset) => {
+                                self.write_output(format!(
+                                    "sub r{} sp {offset}",
+                                    VariableScope::TEMP_STACK_REGISTER
+                                ))?;
+                                self.write_output(format!(
+                                    "get r{} db r{}",
+                                    VariableScope::RETURN_REGISTER,
+                                    VariableScope::TEMP_STACK_REGISTER
+                                ))?;
+                            }
+                            VariableLocation::Device(_) => {
+                                return Err(Error::Unknown(
+                                    "You can not return a device from a function.".into(),
+                                    Some(var_name.span),
+                                ));
+                            }
+                        },
+                        Err(_) => {
+                            self.errors.push(Error::UnknownIdentifier(
+                                var_name.node.clone(),
+                                var_name.span,
+                            ));
+                            // Proceed with dummy
+                        }
                     }
                 }
+                Expression::Literal(spanned_lit) => match spanned_lit.node {
+                    Literal::Number(num) => {
+                        self.emit_variable_assignment(
+                            Cow::from("returnValue"),
+                            &VariableLocation::Persistant(VariableScope::RETURN_REGISTER),
+                            Cow::from(num.to_string()),
+                        )?;
+                    }
+                    Literal::Boolean(b) => {
+                        let val = if b { "1" } else { "0" };
+                        self.emit_variable_assignment(
+                            Cow::from("returnValue"),
+                            &VariableLocation::Persistant(VariableScope::RETURN_REGISTER),
+                            Cow::from(val.to_string()),
+                        )?;
+                    }
+                    _ => {}
+                },
+                Expression::Binary(bin_expr) => {
+                    let result = self.expression_binary(bin_expr, scope)?;
+                    let result_reg = self.resolve_register(&result.location)?;
+                    self.write_output(format!(
+                        "move r{} {}",
+                        VariableScope::RETURN_REGISTER,
+                        result_reg
+                    ))?;
+                    if let Some(name) = result.temp_name {
+                        scope.free_temp(name, None)?;
+                    }
+                }
+                Expression::Logical(log_expr) => {
+                    let result = self.expression_logical(log_expr, scope)?;
+                    let result_reg = self.resolve_register(&result.location)?;
+                    self.write_output(format!(
+                        "move r{} {}",
+                        VariableScope::RETURN_REGISTER,
+                        result_reg
+                    ))?;
+                    if let Some(name) = result.temp_name {
+                        scope.free_temp(name, None)?;
+                    }
+                }
+                Expression::MemberAccess(access) => {
+                    // Return result of member access
+                    let res_opt = self.expression(
+                        Spanned {
+                            node: Expression::MemberAccess(access),
+                            span: expr.span,
+                        },
+                        scope,
+                    )?;
+                    if let Some(res) = res_opt {
+                        let reg = self.resolve_register(&res.location)?;
+                        self.write_output(format!(
+                            "move r{} {}",
+                            VariableScope::RETURN_REGISTER,
+                            reg
+                        ))?;
+                        if let Some(temp) = res.temp_name {
+                            scope.free_temp(temp, None)?;
+                        }
+                    }
+                }
+                _ => {
+                    return Err(Error::Unknown(
+                        format!("Unsupported `return` statement: {:?}", expr),
+                        None,
+                    ));
+                }
             }
-            _ => {
+        }
-                return Err(Error::Unknown(
+
-                    format!("Unsupported `return` statement: {:?}", expr),
+        if let Some(label) = &self.current_return_label {
-                    None,
+            self.write_output(format!("j {}", label))?;
-                ));
+        } else {
-            }
+            return Err(Error::Unknown(
+                "Return statement used outside of function context.".into(),
+                None,
+            ));
         }
 
         Ok(VariableLocation::Persistant(VariableScope::RETURN_REGISTER))
@@ -2344,8 +2363,13 @@ impl<'a, 'w, W: std::io::Write> Compiler<'a, 'w, W> {
         }
 
         self.write_output("push ra")?;
+
+        let return_label = self.next_label_name();
+
+        let prev_return_label = self.current_return_label.replace(return_label.clone());
+
         block_scope.add_variable(
-            Cow::from(format!("{}_ra", name.node)),
+            return_label.clone(),
             LocationRequest::Stack,
             Some(name.span),
         )?;
@@ -2353,7 +2377,7 @@ impl<'a, 'w, W: std::io::Write> Compiler<'a, 'w, W> {
         for expr in body.0 {
             match expr.node {
                 Expression::Return(ret_expr) => {
-                    self.expression_return(*ret_expr, &mut block_scope)?;
+                    self.expression_return(ret_expr, &mut block_scope)?;
                 }
                 _ => {
                     // Swallow internal errors
@@ -2372,11 +2396,13 @@ impl<'a, 'w, W: std::io::Write> Compiler<'a, 'w, W> {
         }
 
         // Get the saved return address and save it back into `ra`
-        let ra_res =
+        let ra_res = block_scope.get_location_of(&return_label, Some(name.span));
-            block_scope.get_location_of(&Cow::from(format!("{}_ra", name.node)), Some(name.span));
 
         let ra_stack_offset = match ra_res {
-            Ok(VariableLocation::Stack(offset)) => offset,
+            Ok(VariableLocation::Stack(offset)) => {
+                block_scope.free_temp(return_label.clone(), None)?;
+                offset
+            }
             _ => {
                 // If we can't find RA, we can't return properly.
                 // This usually implies a compiler bug or scope tracking error.
@@ -2387,6 +2413,10 @@ impl<'a, 'w, W: std::io::Write> Compiler<'a, 'w, W> {
             }
         };
 
+        self.current_return_label = prev_return_label;
+
+        self.write_output(format!("{}:", return_label))?;
+
         self.write_output(format!(
             "sub r{0} sp {ra_stack_offset}",
             VariableScope::TEMP_STACK_REGISTER

rust_compiler/libs/compiler/src/variable_manager.rs

@@ -94,19 +94,21 @@ impl<'a, 'b> VariableScope<'a, 'b> {
     pub const RETURN_REGISTER: u8 = 15;
     pub const TEMP_STACK_REGISTER: u8 = 0;
 
-    pub fn registers(&self) -> impl Iterator<Item = &u8> {
+    pub fn registers(&self) -> Vec<u8> {
-        self.var_lookup_table
+        let mut used = Vec::new();
-            .values()
+
-            .filter(|val| {
+        for r in TEMP {
-                matches!(
+            if !self.temporary_vars.contains(&r) {
-                    val,
+                used.push(r);
-                    VariableLocation::Temporary(_) | VariableLocation::Persistant(_)
+            }
-                )
+        }
-            })
+
-            .map(|loc| match loc {
+        for r in PERSIST {
-                VariableLocation::Persistant(reg) | VariableLocation::Temporary(reg) => reg,
+            if !self.persistant_vars.contains(&r) {
-                _ => unreachable!(),
+                used.push(r);
-            })
+            }
+        }
+        used
     }
 
     pub fn scoped(parent: &'b VariableScope<'a, 'b>) -> Self {

rust_compiler/libs/parser/src/lib.rs

@@ -1225,18 +1225,34 @@ impl<'a> Parser<'a> {
             // Need to capture return span
             let ret_start_span = Self::token_to_span(&current_token);
             self.assign_next()?;
-            let expression = self.expression()?.ok_or(Error::UnexpectedEOF)?;
+
+            let expr = if token_matches!(
+                self.current_token.as_ref().ok_or(Error::UnexpectedEOF)?,
+                TokenType::Symbol(Symbol::Semicolon)
+            ) {
+                // rewind 1 token so we can check for the semicolon at the bottom of this function.
+                self.tokenizer.seek(SeekFrom::Current(-1))?;
+                None
+            } else {
+                Some(self.expression()?.ok_or(Error::UnexpectedEOF)?)
+            };
 
             let ret_span = Span {
                 start_line: ret_start_span.start_line,
                 start_col: ret_start_span.start_col,
-                end_line: expression.span.end_line,
+                end_line: expr
-                end_col: expression.span.end_col,
+                    .as_ref()
+                    .map(|e| e.span.end_line)
+                    .unwrap_or(ret_start_span.end_line),
+                end_col: expr
+                    .as_ref()
+                    .map(|e| e.span.end_col)
+                    .unwrap_or(ret_start_span.end_col),
             };
 
             let return_expr = Spanned {
                 span: ret_span,
-                node: Expression::Return(boxed!(expression)),
+                node: Expression::Return(expr.map(Box::new)),
             };
             expressions.push(return_expr);
 

rust_compiler/libs/parser/src/tree_node.rs

@@ -381,7 +381,7 @@ pub enum Expression<'a> {
     MethodCall(Spanned<MethodCallExpression<'a>>),
     Negation(Box<Spanned<Expression<'a>>>),
     Priority(Box<Spanned<Expression<'a>>>),
-    Return(Box<Spanned<Expression<'a>>>),
+    Return(Option<Box<Spanned<Expression<'a>>>>),
     Syscall(Spanned<SysCall<'a>>),
     Ternary(Spanned<TernaryExpression<'a>>),
     Variable(Spanned<Cow<'a, str>>),
@@ -409,7 +409,15 @@ impl<'a> std::fmt::Display for Expression<'a> {
             Expression::MethodCall(e) => write!(f, "{}", e),
             Expression::Negation(e) => write!(f, "(-{})", e),
             Expression::Priority(e) => write!(f, "({})", e),
-            Expression::Return(e) => write!(f, "(return {})", e),
+            Expression::Return(e) => write!(
+                f,
+                "(return {})",
+                if let Some(e) = e {
+                    e.to_string()
+                } else {
+                    "".to_string()
+                }
+            ),
             Expression::Syscall(e) => write!(f, "{}", e),
             Expression::Ternary(e) => write!(f, "{}", e),
             Expression::Variable(id) => write!(f, "{}", id),