From d6d8db272aa118fafe7f43ab74efa2aa65cbb9df Mon Sep 17 00:00:00 2001 From: Morgan Thomas Date: Thu, 11 Dec 2025 15:52:20 +0000 Subject: [PATCH 1/2] check for wrong number of returned vars in function calls --- crates/lean_compiler/src/a_simplify_lang.rs | 17 +++++++++++++++++ crates/lean_compiler/tests/test_compiler.rs | 15 +++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/crates/lean_compiler/src/a_simplify_lang.rs b/crates/lean_compiler/src/a_simplify_lang.rs index 4c86348d..56ecee30 100644 --- a/crates/lean_compiler/src/a_simplify_lang.rs +++ b/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, lines: &[Line], counters: &mut Counters, new_functions: &mut BTreeMap, @@ -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,16 @@ 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)) diff --git a/crates/lean_compiler/tests/test_compiler.rs b/crates/lean_compiler/tests/test_compiler.rs index 2267bbe9..de2c8b49 100644 --- a/crates/lean_compiler/tests/test_compiler.rs +++ b/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) From 3e90c263c1f3736a644030276026b280afb50903 Mon Sep 17 00:00:00 2001 From: Morgan Thomas Date: Thu, 11 Dec 2025 16:10:39 +0000 Subject: [PATCH 2/2] check for wrong number of returned vars in return statements --- crates/lean_compiler/src/a_simplify_lang.rs | 13 +++++++++++++ crates/lean_compiler/tests/test_compiler.rs | 17 ++++++++++++++++- 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/crates/lean_compiler/src/a_simplify_lang.rs b/crates/lean_compiler/src/a_simplify_lang.rs index 56ecee30..f985a3cd 100644 --- a/crates/lean_compiler/src/a_simplify_lang.rs +++ b/crates/lean_compiler/src/a_simplify_lang.rs @@ -161,6 +161,7 @@ pub fn simplify_program(mut program: Program) -> SimpleProgram { let mut array_manager = ArrayManager::default(); let simplified_instructions = simplify_lines( &program.functions, + func.n_returned_vars, &func.body, &mut counters, &mut new_functions, @@ -436,8 +437,10 @@ impl ArrayManager { } } +#[allow(clippy::too_many_arguments)] fn simplify_lines( functions: &BTreeMap, + n_returned_vars: usize, lines: &[Line], counters: &mut Counters, new_functions: &mut BTreeMap, @@ -458,6 +461,7 @@ fn simplify_lines( assert_eq!(*pattern, i, "match patterns should be consecutive, starting from 0"); simple_arms.push(simplify_lines( functions, + n_returned_vars, statements, counters, new_functions, @@ -610,6 +614,7 @@ fn simplify_lines( let mut array_manager_then = array_manager.clone(); let then_branch_simplified = simplify_lines( functions, + n_returned_vars, then_branch, counters, new_functions, @@ -622,6 +627,7 @@ fn simplify_lines( let else_branch_simplified = simplify_lines( functions, + n_returned_vars, else_branch, counters, new_functions, @@ -672,6 +678,7 @@ fn simplify_lines( replace_vars_for_unroll(&mut body_copy, iterator, unroll_index, i, &internal_variables); unrolled_lines.extend(simplify_lines( functions, + 0, &body_copy, counters, new_functions, @@ -696,6 +703,7 @@ fn simplify_lines( array_manager.valid.clear(); let simplified_body = simplify_lines( functions, + 0, body, counters, new_functions, @@ -793,6 +801,11 @@ fn simplify_lines( } Line::FunctionRet { return_data } => { assert!(!in_a_loop, "Function return inside a loop is not currently supported"); + assert!( + return_data.len() == n_returned_vars, + "Wrong number of return values in return statement; expected {n_returned_vars} but got {}", + return_data.len() + ); let simplified_return_data = return_data .iter() .map(|ret| simplify_expr(ret, &mut res, counters, array_manager, const_malloc)) diff --git a/crates/lean_compiler/tests/test_compiler.rs b/crates/lean_compiler/tests/test_compiler.rs index de2c8b49..3142c6d7 100644 --- a/crates/lean_compiler/tests/test_compiler.rs +++ b/crates/lean_compiler/tests/test_compiler.rs @@ -40,7 +40,7 @@ fn test_duplicate_constant_name() { #[test] #[should_panic] -fn test_wrong_n_returned_vars() { +fn test_wrong_n_returned_vars_1() { let program = r#" fn main() { a, b = f(); @@ -53,6 +53,21 @@ fn test_wrong_n_returned_vars() { compile_and_run(program.to_string(), (&[], &[]), DEFAULT_NO_VEC_RUNTIME_MEMORY, false); } +#[test] +#[should_panic] +fn test_wrong_n_returned_vars_2() { + let program = r#" + fn main() { + a = f(); + } + + fn f() -> 1 { + return 0, 1; + } + "#; + 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)