check for wrong number of returned vars in function calls

Author: morganthomas

crates/lean_compiler/src/a_simplify_lang.rs

@@ -160,6 +160,7 @@ pub fn simplify_program(mut program: Program) -> SimpleProgram {
     for (name, func) in &program.functions {
         let mut array_manager = ArrayManager::default();
         let simplified_instructions = simplify_lines(
+            &program.functions,
             &func.body,
             &mut counters,
             &mut new_functions,
@@ -436,6 +437,7 @@ impl ArrayManager {
 }
 
 fn simplify_lines(
+    functions: &BTreeMap<String, Function>,
     lines: &[Line],
     counters: &mut Counters,
     new_functions: &mut BTreeMap<String, SimpleFunction>,
@@ -455,6 +457,7 @@ fn simplify_lines(
                 for (i, (pattern, statements)) in arms.iter().enumerate() {
                     assert_eq!(*pattern, i, "match patterns should be consecutive, starting from 0");
                     simple_arms.push(simplify_lines(
+                        functions,
                         statements,
                         counters,
                         new_functions,
@@ -606,6 +609,7 @@ fn simplify_lines(
 
                 let mut array_manager_then = array_manager.clone();
                 let then_branch_simplified = simplify_lines(
+                    functions,
                     then_branch,
                     counters,
                     new_functions,
@@ -617,6 +621,7 @@ fn simplify_lines(
                 array_manager_else.valid = array_manager.valid.clone(); // Crucial: remove the access added in the IF branch
 
                 let else_branch_simplified = simplify_lines(
+                    functions,
                     else_branch,
                     counters,
                     new_functions,
@@ -666,6 +671,7 @@ fn simplify_lines(
                         let mut body_copy = body.clone();
                         replace_vars_for_unroll(&mut body_copy, iterator, unroll_index, i, &internal_variables);
                         unrolled_lines.extend(simplify_lines(
+                            functions,
                             &body_copy,
                             counters,
                             new_functions,
@@ -689,6 +695,7 @@ fn simplify_lines(
                 let valid_aux_vars_in_array_manager_before = array_manager.valid.clone();
                 array_manager.valid.clear();
                 let simplified_body = simplify_lines(
+                    functions,
                     body,
                     counters,
                     new_functions,
@@ -763,6 +770,11 @@ fn simplify_lines(
                 return_data,
                 line_number,
             } => {
+                let function = functions.get(function_name).expect("Function used but not defined: {function_name}");
+                if return_data.len() != function.n_returned_vars {
+                    panic!("Expected {} returned vars in call to {function_name}", function.n_returned_vars);
+                }
+
                 let simplified_args = args
                     .iter()
                     .map(|arg| simplify_expr(arg, &mut res, counters, array_manager, const_malloc))

crates/lean_compiler/tests/test_compiler.rs

@@ -38,6 +38,21 @@ fn test_duplicate_constant_name() {
     compile_and_run(program.to_string(), (&[], &[]), DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
 
+#[test]
+#[should_panic]
+fn test_wrong_n_returned_vars() {
+    let program = r#"
+    fn main() {
+        a, b = f();
+    }
+
+    fn f() -> 1 {
+        return 0;
+    }
+    "#;
+    compile_and_run(program.to_string(), (&[], &[]), DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
+}
+
 #[test]
 fn test_fibonacci_program() {
     // a program to check the value of the 30th Fibonacci number (832040)