wip: issue/88

morganthomas · morganthomas · commit f3942d28d625 · 2025-11-30T01:50:58.000Z
diff --git a/crates/lean_compiler/src/a_simplify_lang.rs b/crates/lean_compiler/src/a_simplify_lang.rs
@@ -371,7 +371,6 @@ struct ArrayManager {
 pub struct ConstMalloc {
     counter: usize,
     map: BTreeMap<Var, ConstMallocLabel>,
-    forbidden_vars: BTreeSet<Var>, // vars shared between branches of an if/else
 }
 
 impl ArrayManager {
@@ -397,8 +396,8 @@ fn simplify_lines(
     let mut res = Vec::new();
     for line in lines {
         match line {
-            Line::ForwardDeclaration { var: _ } => {
-                todo!();
+            Line::ForwardDeclaration { var } => {
+                res.push(SimpleLine::ForwardDeclaration { var: var.clone() });
             },
             Line::Match { value, arms } => {
                 let simple_value = simplify_expr(value, &mut res, counters, array_manager, const_malloc);
@@ -555,17 +554,6 @@ fn simplify_lines(
                     }
                 };
 
-                let forbidden_vars_before = const_malloc.forbidden_vars.clone();
-
-                let then_internal_vars = find_variable_usage(then_branch).0;
-                let else_internal_vars = find_variable_usage(else_branch).0;
-                let new_forbidden_vars = then_internal_vars
-                    .intersection(&else_internal_vars)
-                    .cloned()
-                    .collect::<BTreeSet<_>>();
-
-                const_malloc.forbidden_vars.extend(new_forbidden_vars);
-
                 let mut array_manager_then = array_manager.clone();
                 let then_branch_simplified = simplify_lines(
                     then_branch,
@@ -587,8 +575,6 @@ fn simplify_lines(
                     const_malloc,
                 );
 
-                const_malloc.forbidden_vars = forbidden_vars_before;
-
                 *array_manager = array_manager_else.clone();
                 // keep the intersection both branches
                 array_manager.valid = array_manager
@@ -650,6 +636,8 @@ fn simplify_lines(
                     counter: const_malloc.counter,
                     ..ConstMalloc::default()
                 };
+                // TODO: what is array manager, and does it need to be updated
+                // to make block-level scoping work?
                 let valid_aux_vars_in_array_manager_before = array_manager.valid.clone();
                 array_manager.valid.clear();
                 let simplified_body = simplify_lines(
@@ -781,12 +769,8 @@ fn simplify_lines(
                 let simplified_size = simplify_expr(size, &mut res, counters, array_manager, const_malloc);
                 let simplified_vectorized_len =
                     simplify_expr(vectorized_len, &mut res, counters, array_manager, const_malloc);
-                if simplified_size.is_constant() && !*vectorized && const_malloc.forbidden_vars.contains(var) {
-                    println!("TODO: Optimization missed: Requires to align const malloc in if/else branches");
-                }
                 match simplified_size {
-                    SimpleExpr::Constant(const_size) if !*vectorized && !const_malloc.forbidden_vars.contains(var) => {
-                        // TODO do this optimization even if we are in an if/else branch
+                    SimpleExpr::Constant(const_size) if !*vectorized => {
                         let label = const_malloc.counter;
                         const_malloc.counter += 1;
                         const_malloc.map.insert(var.clone(), label);
@@ -807,7 +791,6 @@ fn simplify_lines(
                 }
             }
             Line::DecomposeBits { var, to_decompose } => {
-                assert!(!const_malloc.forbidden_vars.contains(var), "TODO");
                 let simplified_to_decompose = to_decompose
                     .iter()
                     .map(|expr| simplify_expr(expr, &mut res, counters, array_manager, const_malloc))
@@ -822,7 +805,6 @@ fn simplify_lines(
                 });
             }
             Line::DecomposeCustom { var, to_decompose } => {
-                assert!(!const_malloc.forbidden_vars.contains(var), "TODO");
                 let simplified_to_decompose = to_decompose
                     .iter()
                     .map(|expr| simplify_expr(expr, &mut res, counters, array_manager, const_malloc))
@@ -949,9 +931,8 @@ pub fn find_variable_usage(lines: &[Line]) -> (BTreeSet<Var>, BTreeSet<Var>) {
 
     for line in lines {
         match line {
-            Line::ForwardDeclaration { var: _ } => {
-                todo!();
-                // internal_vars.push(var.clone()); // if declaring is using; otherwise do nothing
+            Line::ForwardDeclaration { var } => {
+                internal_vars.insert(var.clone());
             },
             Line::Match { value, arms } => {
                 on_new_expr(value, &internal_vars, &mut external_vars);
@@ -1997,6 +1978,9 @@ impl SimpleLine {
     fn to_string_with_indent(&self, indent: usize) -> String {
         let spaces = "    ".repeat(indent);
         let line_str = match self {
+            Self::ForwardDeclaration { var } => {
+                format!("var {var}")
+            }
             Self::Match { value, arms } => {
                 let arms_str = arms
                     .iter()
diff --git a/crates/lean_compiler/src/b_compile_intermediate.rs b/crates/lean_compiler/src/b_compile_intermediate.rs
@@ -19,6 +19,7 @@ struct Compiler {
     stack_size: usize,
 }
 
+#[derive(Default)]
 struct StackFrameLayout {
     // Innermost lexical scope last
     scopes: Vec<ScopeLayout>,
@@ -36,7 +37,7 @@ impl Compiler {
             VarOrConstMallocAccess::Var(var) => {
                 for scope in self.stack_frame_layout.scopes.iter().rev() {
                     if let Some(offset) = scope.var_positions.get(var) {
-                        return offset.into();
+                        return (*offset).into();
                     }
                 }
                 panic!("Variable {var} not in scope");
@@ -45,9 +46,9 @@ impl Compiler {
                 for scope in self.stack_frame_layout.scopes.iter().rev() {
                     if let Some(base) = scope.const_mallocs.get(malloc_label) {
                         return ConstExpression::Binary {
-                            left: Box::new(base.into()),
+                            left: Box::new((*base).into()),
                             operation: HighLevelOperation::Add,
-                            right: Box::new(offset.clone()),
+                            right: Box::new((*offset).clone()),
                         };
                     }
                 }
@@ -82,7 +83,10 @@ impl IntermediateValue {
             },
             SimpleExpr::Constant(c) => Self::Constant(c.clone()),
             SimpleExpr::ConstMallocAccess { malloc_label, offset } => Self::MemoryAfterFp {
-                offset: compiler.get_offset(SimpleExpr::ConstMallocAccess { malloc_label, offset }),
+                offset: compiler.get_offset(&VarOrConstMallocAccess::ConstMallocAccess {
+                    malloc_label: *malloc_label,
+                    offset: offset.clone()
+                }),
             },
         }
     }
@@ -115,15 +119,14 @@ fn compile_function(
 ) -> Result<Vec<IntermediateInstruction>, String> {
     // memory layout: pc, fp, args, return_vars, internal_vars
     let mut stack_pos = 2; // Reserve space for pc and fp
-    let mut var_positions = BTreeMap::new();
-    let mut function_scope = Scope::default();
+    let mut function_scope_layout = ScopeLayout::default();
     compiler.stack_frame_layout = StackFrameLayout {
-        scopes: vec![function_scope],
+        scopes: vec![function_scope_layout],
     };
-    let function_scope = &mut compiler.stack_frame_layout.scopes[0];
+    let function_scope_layout = &mut compiler.stack_frame_layout.scopes[0];
 
     for (i, var) in function.arguments.iter().enumerate() {
-        function_scope.var_positions.insert(var.clone(), stack_pos + i);
+        function_scope_layout.var_positions.insert(var.clone(), stack_pos + i);
     }
     stack_pos += function.arguments.len();
 
@@ -150,14 +153,14 @@ fn compile_lines(
     final_jump: Option<Label>,
     declared_vars: &mut BTreeSet<Var>,
 ) -> Result<Vec<IntermediateInstruction>, String> {
-    // TODO: amend stack frame layout with internal variables.
-
     let mut instructions = Vec::new();
 
     for (i, line) in lines.iter().enumerate() {
         match line {
-            SimpleLine::ForwardDeclaration => {
-                todo!(); // amend stack frame layout
+            SimpleLine::ForwardDeclaration { var } => {
+                let mut current_scope_layout = compiler.stack_frame_layout.scopes.last_mut().unwrap();
+                current_scope_layout.var_positions.insert(var.clone(), compiler.stack_size);
+                compiler.stack_size += 1;
             }
 
             SimpleLine::Assignment {
@@ -166,18 +169,23 @@ fn compile_lines(
                 arg0,
                 arg1,
             } => {
-                todo!(); // amend stack frame layout
-                instructions.push(IntermediateInstruction::computation(
-                    *operation,
-                    IntermediateValue::from_simple_expr(arg0, compiler),
-                    IntermediateValue::from_simple_expr(arg1, compiler),
-                    IntermediateValue::from_var_or_const_malloc_access(var, compiler),
-                ));
-
                 mark_vars_as_declared(&[arg0, arg1], declared_vars);
+                let arg0 = IntermediateValue::from_simple_expr(arg0, compiler);
+                let arg1 = IntermediateValue::from_simple_expr(arg1, compiler);
+
                 if let VarOrConstMallocAccess::Var(var) = var {
                     declared_vars.insert(var.clone());
+                    let mut current_scope_layout = compiler.stack_frame_layout.scopes.last_mut().unwrap();
+                    current_scope_layout.var_positions.insert(var.clone(), compiler.stack_size);
+                    compiler.stack_size += 1;
                 }
+
+                instructions.push(IntermediateInstruction::computation(
+                    *operation,
+                    arg0,
+                    arg1,
+                    IntermediateValue::from_var_or_const_malloc_access(var, compiler),
+                ));
             }
 
             SimpleLine::TestZero { operation, arg0, arg1 } => {
@@ -192,7 +200,9 @@ fn compile_lines(
             }
 
             SimpleLine::Match { value, arms } => {
-                todo!(); // amend stack frame layout
+                let saved_stack_size = compiler.stack_size;
+                compiler.stack_frame_layout.scopes.push(ScopeLayout::default());
+
                 let match_index = compiler.match_blocks.len();
                 let end_label = Label::match_end(match_index);
 
@@ -254,6 +264,9 @@ fn compile_lines(
                 let remaining = compile_lines(function_name, &lines[i + 1..], compiler, final_jump, declared_vars)?;
                 compiler.bytecode.insert(end_label, remaining);
 
+                compiler.stack_frame_layout.scopes.pop();
+                compiler.stack_size = saved_stack_size;
+
                 return Ok(instructions);
             }
 
@@ -263,7 +276,9 @@ fn compile_lines(
                 else_branch,
                 line_number,
             } => {
-                todo!(); // amend stack frame layout
+                let saved_stack_size = compiler.stack_size;
+                compiler.stack_frame_layout.scopes.push(ScopeLayout::default());
+
                 validate_vars_declared(&[condition], declared_vars)?;
 
                 let if_id = compiler.if_counter;
@@ -368,6 +383,9 @@ fn compile_lines(
                 let remaining = compile_lines(function_name, &lines[i + 1..], compiler, final_jump, declared_vars)?;
                 compiler.bytecode.insert(end_label, remaining);
 
+                compiler.stack_frame_layout.scopes.pop();
+                compiler.stack_size = saved_stack_size;
+
                 return Ok(instructions);
             }
 
@@ -393,7 +411,6 @@ fn compile_lines(
                 return_data,
                 line_number,
             } => {
-                todo!(); // amend stack frame layout
                 let call_id = compiler.call_counter;
                 compiler.call_counter += 1;
                 let return_label = Label::return_from_call(call_id, *line_number);
@@ -411,6 +428,11 @@ fn compile_lines(
 
                 validate_vars_declared(args, declared_vars)?;
                 declared_vars.extend(return_data.iter().cloned());
+                for var in return_data.iter() {
+                    let mut current_scope_layout = compiler.stack_frame_layout.scopes.last_mut().unwrap();
+                    current_scope_layout.var_positions.insert(var.clone(), compiler.stack_size);
+                    compiler.stack_size += 1;
+                }
 
                 let after_call = {
                     let mut instructions = Vec::new();
@@ -443,7 +465,6 @@ fn compile_lines(
             }
 
             SimpleLine::Precompile { table, args, .. } => {
-                todo!(); // amend stack frame layout?
                 if *table == Table::poseidon24() {
                     assert_eq!(args.len(), 3);
                 } else {
@@ -486,7 +507,9 @@ fn compile_lines(
                 vectorized,
                 vectorized_len,
             } => {
-                todo!(); // amend stack frame layout
+                let mut current_scope_layout = compiler.stack_frame_layout.scopes.last_mut().unwrap();
+                current_scope_layout.var_positions.insert(var.clone(), compiler.stack_size);
+                compiler.stack_size += 1;
                 declared_vars.insert(var.clone());
                 instructions.push(IntermediateInstruction::RequestMemory {
                     offset: compiler.get_offset(&var.clone().into()),
@@ -496,7 +519,6 @@ fn compile_lines(
                 });
             }
             SimpleLine::ConstMalloc { var, size, label } => {
-                todo!(); // amend stack frame layout
                 let size = size.naive_eval().unwrap().to_usize(); // TODO not very good;
                 handle_const_malloc(declared_vars, &mut instructions, compiler, var, size, label);
             }
@@ -505,7 +527,6 @@ fn compile_lines(
                 to_decompose,
                 label,
             } => {
-                todo!(); // amend stack frame layout
                 instructions.push(IntermediateInstruction::DecomposeBits {
                     res_offset: compiler.stack_size,
                     to_decompose: to_decompose
@@ -528,7 +549,6 @@ fn compile_lines(
                 to_decompose,
                 label,
             } => {
-                todo!(); // amend stack frame layout
                 instructions.push(IntermediateInstruction::DecomposeCustom {
                     res_offset: compiler.stack_size,
                     to_decompose: to_decompose
@@ -547,7 +567,9 @@ fn compile_lines(
                 );
             }
             SimpleLine::CounterHint { var } => {
-                todo!(); // amend stack frame layout
+                let mut current_scope_layout = compiler.stack_frame_layout.scopes.last_mut().unwrap();
+                current_scope_layout.var_positions.insert(var.clone(), compiler.stack_size);
+                compiler.stack_size += 1;
                 declared_vars.insert(var.clone());
                 instructions.push(IntermediateInstruction::CounterHint {
                     res_offset: compiler
@@ -590,7 +612,6 @@ fn handle_const_malloc(
     size: usize,
     label: &ConstMallocLabel,
 ) {
-    todo!(); // amend stack frame layout
     declared_vars.insert(var.clone());
     instructions.push(IntermediateInstruction::Computation {
         operation: Operation::Add,
@@ -600,7 +621,8 @@ fn handle_const_malloc(
             offset: compiler.get_offset(&var.clone().into()),
         },
     });
-    compiler.const_mallocs.insert(*label, compiler.stack_size);
+    let current_scope = compiler.stack_frame_layout.scopes.last_mut().unwrap();
+    current_scope.const_mallocs.insert(*label, compiler.stack_size);
     compiler.stack_size += size;
 }
 
@@ -688,11 +710,13 @@ fn compile_function_ret(
     });
 }
 
-// TODO: delete this?
 fn find_internal_vars(lines: &[SimpleLine]) -> BTreeSet<Var> {
     let mut internal_vars = BTreeSet::new();
     for line in lines {
         match line {
+            SimpleLine::ForwardDeclaration { var } => {
+                internal_vars.insert(var.clone());
+            }
             SimpleLine::Match { arms, .. } => {
                 for arm in arms {
                     internal_vars.extend(find_internal_vars(arm));
diff --git a/crates/lean_compiler/src/lang.rs b/crates/lean_compiler/src/lang.rs
@@ -398,6 +398,7 @@ impl Context {
     }
 }
 
+#[derive(Default)]
 pub struct Scope {
     /// A set of declared variables.
     pub vars: BTreeSet<Var>,
diff --git a/crates/lean_compiler/src/lib.rs b/crates/lean_compiler/src/lib.rs
@@ -16,9 +16,9 @@ pub fn compile_program(program: String) -> Bytecode {
     let (parsed_program, function_locations) = parse_program(&program).unwrap();
     // println!("Parsed program: {}", parsed_program.to_string());
     let simple_program = simplify_program(parsed_program);
-    // println!("Simplified program: {}", simple_program.to_string());
+    println!("Simplified program: {}", simple_program.to_string());
     let intermediate_bytecode = compile_to_intermediate_bytecode(simple_program).unwrap();
-    // println!("Intermediate Bytecode:\n\n{}", intermediate_bytecode.to_string());
+    println!("Intermediate Bytecode:\n\n{}", intermediate_bytecode.to_string());
 
     // println!("Function Locations: \n");
     // for (loc, name) in function_locations.iter() {

Original file line number	Diff line number	Diff line change
`@@ -398,6 +398,7 @@ impl Context {`
`398`	`398`	`}`
`399`	`399`	`}`
`400`	`400`
	`401`	`+#[derive(Default)]`
`401`	`402`	`pub struct Scope {`
`402`	`403`	`/// A set of declared variables.`
`403`	`404`	`pub vars: BTreeSet<Var>,`