Skip to content

Commit 7b0ef06

Browse files
morganthomasmorganthomas
committed
wip: issue/88
1 parent 597f86f commit 7b0ef06

File tree

3 files changed

+56
-49
lines changed

3 files changed

+56
-49
lines changed

‎crates/lean_compiler/src/a_simplify_lang.rs‎

Lines changed: 8 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -371,7 +371,6 @@ struct ArrayManager {
371371
pub struct ConstMalloc {
372372
counter: usize,
373373
map: BTreeMap<Var, ConstMallocLabel>,
374-
forbidden_vars: BTreeSet<Var>, // vars shared between branches of an if/else
375374
}
376375

377376
impl ArrayManager {
@@ -397,8 +396,8 @@ fn simplify_lines(
397396
let mut res = Vec::new();
398397
for line in lines {
399398
match line {
400-
Line::ForwardDeclaration { var: _ } => {
401-
todo!();
399+
Line::ForwardDeclaration { var } => {
400+
res.push(SimpleLine::ForwardDeclaration { var: var.clone() });
402401
},
403402
Line::Match { value, arms } => {
404403
let simple_value = simplify_expr(value, &mut res, counters, array_manager, const_malloc);
@@ -555,17 +554,6 @@ fn simplify_lines(
555554
}
556555
};
557556

558-
let forbidden_vars_before = const_malloc.forbidden_vars.clone();
559-
560-
let then_internal_vars = find_variable_usage(then_branch).0;
561-
let else_internal_vars = find_variable_usage(else_branch).0;
562-
let new_forbidden_vars = then_internal_vars
563-
.intersection(&else_internal_vars)
564-
.cloned()
565-
.collect::<BTreeSet<_>>();
566-
567-
const_malloc.forbidden_vars.extend(new_forbidden_vars);
568-
569557
let mut array_manager_then = array_manager.clone();
570558
let then_branch_simplified = simplify_lines(
571559
then_branch,
@@ -587,8 +575,6 @@ fn simplify_lines(
587575
const_malloc,
588576
);
589577

590-
const_malloc.forbidden_vars = forbidden_vars_before;
591-
592578
*array_manager = array_manager_else.clone();
593579
// keep the intersection both branches
594580
array_manager.valid = array_manager
@@ -781,12 +767,8 @@ fn simplify_lines(
781767
let simplified_size = simplify_expr(size, &mut res, counters, array_manager, const_malloc);
782768
let simplified_vectorized_len =
783769
simplify_expr(vectorized_len, &mut res, counters, array_manager, const_malloc);
784-
if simplified_size.is_constant() && !*vectorized && const_malloc.forbidden_vars.contains(var) {
785-
println!("TODO: Optimization missed: Requires to align const malloc in if/else branches");
786-
}
787770
match simplified_size {
788-
SimpleExpr::Constant(const_size) if !*vectorized && !const_malloc.forbidden_vars.contains(var) => {
789-
// TODO do this optimization even if we are in an if/else branch
771+
SimpleExpr::Constant(const_size) if !*vectorized => {
790772
let label = const_malloc.counter;
791773
const_malloc.counter += 1;
792774
const_malloc.map.insert(var.clone(), label);
@@ -807,7 +789,6 @@ fn simplify_lines(
807789
}
808790
}
809791
Line::DecomposeBits { var, to_decompose } => {
810-
assert!(!const_malloc.forbidden_vars.contains(var), "TODO");
811792
let simplified_to_decompose = to_decompose
812793
.iter()
813794
.map(|expr| simplify_expr(expr, &mut res, counters, array_manager, const_malloc))
@@ -822,7 +803,6 @@ fn simplify_lines(
822803
});
823804
}
824805
Line::DecomposeCustom { var, to_decompose } => {
825-
assert!(!const_malloc.forbidden_vars.contains(var), "TODO");
826806
let simplified_to_decompose = to_decompose
827807
.iter()
828808
.map(|expr| simplify_expr(expr, &mut res, counters, array_manager, const_malloc))
@@ -949,9 +929,8 @@ pub fn find_variable_usage(lines: &[Line]) -> (BTreeSet<Var>, BTreeSet<Var>) {
949929

950930
for line in lines {
951931
match line {
952-
Line::ForwardDeclaration { var: _ } => {
953-
todo!();
954-
// internal_vars.push(var.clone()); // if declaring is using; otherwise do nothing
932+
Line::ForwardDeclaration { var } => {
933+
internal_vars.insert(var.clone());
955934
},
956935
Line::Match { value, arms } => {
957936
on_new_expr(value, &internal_vars, &mut external_vars);
@@ -1997,6 +1976,9 @@ impl SimpleLine {
19971976
fn to_string_with_indent(&self, indent: usize) -> String {
19981977
let spaces = " ".repeat(indent);
19991978
let line_str = match self {
1979+
Self::ForwardDeclaration { var } => {
1980+
format!("var {var}")
1981+
}
20001982
Self::Match { value, arms } => {
20011983
let arms_str = arms
20021984
.iter()

‎crates/lean_compiler/src/b_compile_intermediate.rs‎

Lines changed: 47 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ struct Compiler {
1919
stack_size: usize,
2020
}
2121

22+
#[derive(Default)]
2223
struct StackFrameLayout {
2324
// Innermost lexical scope last
2425
scopes: Vec<ScopeLayout>,
@@ -36,7 +37,7 @@ impl Compiler {
3637
VarOrConstMallocAccess::Var(var) => {
3738
for scope in self.stack_frame_layout.scopes.iter().rev() {
3839
if let Some(offset) = scope.var_positions.get(var) {
39-
return offset.into();
40+
return (*offset).into();
4041
}
4142
}
4243
panic!("Variable {var} not in scope");
@@ -45,9 +46,9 @@ impl Compiler {
4546
for scope in self.stack_frame_layout.scopes.iter().rev() {
4647
if let Some(base) = scope.const_mallocs.get(malloc_label) {
4748
return ConstExpression::Binary {
48-
left: Box::new(base.into()),
49+
left: Box::new((*base).into()),
4950
operation: HighLevelOperation::Add,
50-
right: Box::new(offset.clone()),
51+
right: Box::new((*offset).clone()),
5152
};
5253
}
5354
}
@@ -82,7 +83,10 @@ impl IntermediateValue {
8283
},
8384
SimpleExpr::Constant(c) => Self::Constant(c.clone()),
8485
SimpleExpr::ConstMallocAccess { malloc_label, offset } => Self::MemoryAfterFp {
85-
offset: compiler.get_offset(SimpleExpr::ConstMallocAccess { malloc_label, offset }),
86+
offset: compiler.get_offset(&VarOrConstMallocAccess::ConstMallocAccess {
87+
malloc_label: *malloc_label,
88+
offset: offset.clone()
89+
}),
8690
},
8791
}
8892
}
@@ -115,15 +119,14 @@ fn compile_function(
115119
) -> Result<Vec<IntermediateInstruction>, String> {
116120
// memory layout: pc, fp, args, return_vars, internal_vars
117121
let mut stack_pos = 2; // Reserve space for pc and fp
118-
let mut var_positions = BTreeMap::new();
119-
let mut function_scope = Scope::default();
122+
let mut function_scope_layout = ScopeLayout::default();
120123
compiler.stack_frame_layout = StackFrameLayout {
121-
scopes: vec![function_scope],
124+
scopes: vec![function_scope_layout],
122125
};
123-
let function_scope = &mut compiler.stack_frame_layout.scopes[0];
126+
let function_scope_layout = &mut compiler.stack_frame_layout.scopes[0];
124127

125128
for (i, var) in function.arguments.iter().enumerate() {
126-
function_scope.var_positions.insert(var.clone(), stack_pos + i);
129+
function_scope_layout.var_positions.insert(var.clone(), stack_pos + i);
127130
}
128131
stack_pos += function.arguments.len();
129132

@@ -156,8 +159,10 @@ fn compile_lines(
156159

157160
for (i, line) in lines.iter().enumerate() {
158161
match line {
159-
SimpleLine::ForwardDeclaration => {
160-
todo!(); // amend stack frame layout
162+
SimpleLine::ForwardDeclaration { var } => {
163+
let mut current_scope_layout = compiler.stack_frame_layout.scopes.last_mut().unwrap();
164+
current_scope_layout.var_positions.insert(var.clone(), compiler.stack_size);
165+
compiler.stack_size += 1;
161166
}
162167

163168
SimpleLine::Assignment {
@@ -166,7 +171,6 @@ fn compile_lines(
166171
arg0,
167172
arg1,
168173
} => {
169-
todo!(); // amend stack frame layout
170174
instructions.push(IntermediateInstruction::computation(
171175
*operation,
172176
IntermediateValue::from_simple_expr(arg0, compiler),
@@ -177,6 +181,9 @@ fn compile_lines(
177181
mark_vars_as_declared(&[arg0, arg1], declared_vars);
178182
if let VarOrConstMallocAccess::Var(var) = var {
179183
declared_vars.insert(var.clone());
184+
let mut current_scope_layout = compiler.stack_frame_layout.scopes.last_mut().unwrap();
185+
current_scope_layout.var_positions.insert(var.clone(), compiler.stack_size);
186+
compiler.stack_size += 1;
180187
}
181188
}
182189

@@ -192,7 +199,9 @@ fn compile_lines(
192199
}
193200

194201
SimpleLine::Match { value, arms } => {
195-
todo!(); // amend stack frame layout
202+
let saved_stack_size = compiler.stack_size;
203+
compiler.stack_frame_layout.scopes.push(ScopeLayout::default());
204+
196205
let match_index = compiler.match_blocks.len();
197206
let end_label = Label::match_end(match_index);
198207

@@ -254,6 +263,9 @@ fn compile_lines(
254263
let remaining = compile_lines(function_name, &lines[i + 1..], compiler, final_jump, declared_vars)?;
255264
compiler.bytecode.insert(end_label, remaining);
256265

266+
compiler.stack_frame_layout.scopes.pop();
267+
compiler.stack_size = saved_stack_size;
268+
257269
return Ok(instructions);
258270
}
259271

@@ -263,7 +275,9 @@ fn compile_lines(
263275
else_branch,
264276
line_number,
265277
} => {
266-
todo!(); // amend stack frame layout
278+
let saved_stack_size = compiler.stack_size;
279+
compiler.stack_frame_layout.scopes.push(ScopeLayout::default());
280+
267281
validate_vars_declared(&[condition], declared_vars)?;
268282

269283
let if_id = compiler.if_counter;
@@ -368,6 +382,9 @@ fn compile_lines(
368382
let remaining = compile_lines(function_name, &lines[i + 1..], compiler, final_jump, declared_vars)?;
369383
compiler.bytecode.insert(end_label, remaining);
370384

385+
compiler.stack_frame_layout.scopes.pop();
386+
compiler.stack_size = saved_stack_size;
387+
371388
return Ok(instructions);
372389
}
373390

@@ -393,7 +410,6 @@ fn compile_lines(
393410
return_data,
394411
line_number,
395412
} => {
396-
todo!(); // amend stack frame layout
397413
let call_id = compiler.call_counter;
398414
compiler.call_counter += 1;
399415
let return_label = Label::return_from_call(call_id, *line_number);
@@ -411,6 +427,11 @@ fn compile_lines(
411427

412428
validate_vars_declared(args, declared_vars)?;
413429
declared_vars.extend(return_data.iter().cloned());
430+
for var in return_data.iter() {
431+
let mut current_scope_layout = compiler.stack_frame_layout.scopes.last_mut().unwrap();
432+
current_scope_layout.var_positions.insert(var.clone(), compiler.stack_size);
433+
compiler.stack_size += 1;
434+
}
414435

415436
let after_call = {
416437
let mut instructions = Vec::new();
@@ -443,7 +464,6 @@ fn compile_lines(
443464
}
444465

445466
SimpleLine::Precompile { table, args, .. } => {
446-
todo!(); // amend stack frame layout?
447467
if *table == Table::poseidon24() {
448468
assert_eq!(args.len(), 3);
449469
} else {
@@ -486,7 +506,9 @@ fn compile_lines(
486506
vectorized,
487507
vectorized_len,
488508
} => {
489-
todo!(); // amend stack frame layout
509+
let mut current_scope_layout = compiler.stack_frame_layout.scopes.last_mut().unwrap();
510+
current_scope_layout.var_positions.insert(var.clone(), compiler.stack_size);
511+
compiler.stack_size += 1;
490512
declared_vars.insert(var.clone());
491513
instructions.push(IntermediateInstruction::RequestMemory {
492514
offset: compiler.get_offset(&var.clone().into()),
@@ -496,7 +518,6 @@ fn compile_lines(
496518
});
497519
}
498520
SimpleLine::ConstMalloc { var, size, label } => {
499-
todo!(); // amend stack frame layout
500521
let size = size.naive_eval().unwrap().to_usize(); // TODO not very good;
501522
handle_const_malloc(declared_vars, &mut instructions, compiler, var, size, label);
502523
}
@@ -505,7 +526,6 @@ fn compile_lines(
505526
to_decompose,
506527
label,
507528
} => {
508-
todo!(); // amend stack frame layout
509529
instructions.push(IntermediateInstruction::DecomposeBits {
510530
res_offset: compiler.stack_size,
511531
to_decompose: to_decompose
@@ -528,7 +548,6 @@ fn compile_lines(
528548
to_decompose,
529549
label,
530550
} => {
531-
todo!(); // amend stack frame layout
532551
instructions.push(IntermediateInstruction::DecomposeCustom {
533552
res_offset: compiler.stack_size,
534553
to_decompose: to_decompose
@@ -547,7 +566,9 @@ fn compile_lines(
547566
);
548567
}
549568
SimpleLine::CounterHint { var } => {
550-
todo!(); // amend stack frame layout
569+
let mut current_scope_layout = compiler.stack_frame_layout.scopes.last_mut().unwrap();
570+
current_scope_layout.var_positions.insert(var.clone(), compiler.stack_size);
571+
compiler.stack_size += 1;
551572
declared_vars.insert(var.clone());
552573
instructions.push(IntermediateInstruction::CounterHint {
553574
res_offset: compiler
@@ -590,7 +611,6 @@ fn handle_const_malloc(
590611
size: usize,
591612
label: &ConstMallocLabel,
592613
) {
593-
todo!(); // amend stack frame layout
594614
declared_vars.insert(var.clone());
595615
instructions.push(IntermediateInstruction::Computation {
596616
operation: Operation::Add,
@@ -600,7 +620,8 @@ fn handle_const_malloc(
600620
offset: compiler.get_offset(&var.clone().into()),
601621
},
602622
});
603-
compiler.const_mallocs.insert(*label, compiler.stack_size);
623+
let current_scope = compiler.stack_frame_layout.scopes.last_mut().unwrap();
624+
current_scope.const_mallocs.insert(*label, compiler.stack_size);
604625
compiler.stack_size += size;
605626
}
606627

@@ -693,6 +714,9 @@ fn find_internal_vars(lines: &[SimpleLine]) -> BTreeSet<Var> {
693714
let mut internal_vars = BTreeSet::new();
694715
for line in lines {
695716
match line {
717+
SimpleLine::ForwardDeclaration { var } => {
718+
internal_vars.insert(var.clone());
719+
}
696720
SimpleLine::Match { arms, .. } => {
697721
for arm in arms {
698722
internal_vars.extend(find_internal_vars(arm));

‎crates/lean_compiler/src/lang.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -398,6 +398,7 @@ impl Context {
398398
}
399399
}
400400

401+
#[derive(Default)]
401402
pub struct Scope {
402403
/// A set of declared variables.
403404
pub vars: BTreeSet<Var>,

0 commit comments

Comments
 (0)