Skip to content

Commit de3ac47

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

File tree

4 files changed

+67
-57
lines changed

4 files changed

+67
-57
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: 57 additions & 30 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,18 +171,23 @@ fn compile_lines(
166171
arg0,
167172
arg1,
168173
} => {
169-
todo!(); // amend stack frame layout
170-
instructions.push(IntermediateInstruction::computation(
171-
*operation,
172-
IntermediateValue::from_simple_expr(arg0, compiler),
173-
IntermediateValue::from_simple_expr(arg1, compiler),
174-
IntermediateValue::from_var_or_const_malloc_access(var, compiler),
175-
));
176-
177174
mark_vars_as_declared(&[arg0, arg1], declared_vars);
175+
let arg0 = IntermediateValue::from_simple_expr(arg0, compiler);
176+
let arg1 = IntermediateValue::from_simple_expr(arg1, compiler);
177+
178178
if let VarOrConstMallocAccess::Var(var) = var {
179179
declared_vars.insert(var.clone());
180+
let mut current_scope_layout = compiler.stack_frame_layout.scopes.last_mut().unwrap();
181+
current_scope_layout.var_positions.insert(var.clone(), compiler.stack_size);
182+
compiler.stack_size += 1;
180183
}
184+
185+
instructions.push(IntermediateInstruction::computation(
186+
*operation,
187+
arg0,
188+
arg1,
189+
IntermediateValue::from_var_or_const_malloc_access(var, compiler),
190+
));
181191
}
182192

183193
SimpleLine::TestZero { operation, arg0, arg1 } => {
@@ -192,7 +202,9 @@ fn compile_lines(
192202
}
193203

194204
SimpleLine::Match { value, arms } => {
195-
todo!(); // amend stack frame layout
205+
let saved_stack_size = compiler.stack_size;
206+
compiler.stack_frame_layout.scopes.push(ScopeLayout::default());
207+
196208
let match_index = compiler.match_blocks.len();
197209
let end_label = Label::match_end(match_index);
198210

@@ -254,6 +266,9 @@ fn compile_lines(
254266
let remaining = compile_lines(function_name, &lines[i + 1..], compiler, final_jump, declared_vars)?;
255267
compiler.bytecode.insert(end_label, remaining);
256268

269+
compiler.stack_frame_layout.scopes.pop();
270+
compiler.stack_size = saved_stack_size;
271+
257272
return Ok(instructions);
258273
}
259274

@@ -263,7 +278,9 @@ fn compile_lines(
263278
else_branch,
264279
line_number,
265280
} => {
266-
todo!(); // amend stack frame layout
281+
let saved_stack_size = compiler.stack_size;
282+
compiler.stack_frame_layout.scopes.push(ScopeLayout::default());
283+
267284
validate_vars_declared(&[condition], declared_vars)?;
268285

269286
let if_id = compiler.if_counter;
@@ -368,6 +385,9 @@ fn compile_lines(
368385
let remaining = compile_lines(function_name, &lines[i + 1..], compiler, final_jump, declared_vars)?;
369386
compiler.bytecode.insert(end_label, remaining);
370387

388+
compiler.stack_frame_layout.scopes.pop();
389+
compiler.stack_size = saved_stack_size;
390+
371391
return Ok(instructions);
372392
}
373393

@@ -393,7 +413,6 @@ fn compile_lines(
393413
return_data,
394414
line_number,
395415
} => {
396-
todo!(); // amend stack frame layout
397416
let call_id = compiler.call_counter;
398417
compiler.call_counter += 1;
399418
let return_label = Label::return_from_call(call_id, *line_number);
@@ -411,6 +430,11 @@ fn compile_lines(
411430

412431
validate_vars_declared(args, declared_vars)?;
413432
declared_vars.extend(return_data.iter().cloned());
433+
for var in return_data.iter() {
434+
let mut current_scope_layout = compiler.stack_frame_layout.scopes.last_mut().unwrap();
435+
current_scope_layout.var_positions.insert(var.clone(), compiler.stack_size);
436+
compiler.stack_size += 1;
437+
}
414438

415439
let after_call = {
416440
let mut instructions = Vec::new();
@@ -443,7 +467,6 @@ fn compile_lines(
443467
}
444468

445469
SimpleLine::Precompile { table, args, .. } => {
446-
todo!(); // amend stack frame layout?
447470
if *table == Table::poseidon24() {
448471
assert_eq!(args.len(), 3);
449472
} else {
@@ -486,7 +509,9 @@ fn compile_lines(
486509
vectorized,
487510
vectorized_len,
488511
} => {
489-
todo!(); // amend stack frame layout
512+
let mut current_scope_layout = compiler.stack_frame_layout.scopes.last_mut().unwrap();
513+
current_scope_layout.var_positions.insert(var.clone(), compiler.stack_size);
514+
compiler.stack_size += 1;
490515
declared_vars.insert(var.clone());
491516
instructions.push(IntermediateInstruction::RequestMemory {
492517
offset: compiler.get_offset(&var.clone().into()),
@@ -496,7 +521,6 @@ fn compile_lines(
496521
});
497522
}
498523
SimpleLine::ConstMalloc { var, size, label } => {
499-
todo!(); // amend stack frame layout
500524
let size = size.naive_eval().unwrap().to_usize(); // TODO not very good;
501525
handle_const_malloc(declared_vars, &mut instructions, compiler, var, size, label);
502526
}
@@ -505,7 +529,6 @@ fn compile_lines(
505529
to_decompose,
506530
label,
507531
} => {
508-
todo!(); // amend stack frame layout
509532
instructions.push(IntermediateInstruction::DecomposeBits {
510533
res_offset: compiler.stack_size,
511534
to_decompose: to_decompose
@@ -528,7 +551,6 @@ fn compile_lines(
528551
to_decompose,
529552
label,
530553
} => {
531-
todo!(); // amend stack frame layout
532554
instructions.push(IntermediateInstruction::DecomposeCustom {
533555
res_offset: compiler.stack_size,
534556
to_decompose: to_decompose
@@ -547,7 +569,9 @@ fn compile_lines(
547569
);
548570
}
549571
SimpleLine::CounterHint { var } => {
550-
todo!(); // amend stack frame layout
572+
let mut current_scope_layout = compiler.stack_frame_layout.scopes.last_mut().unwrap();
573+
current_scope_layout.var_positions.insert(var.clone(), compiler.stack_size);
574+
compiler.stack_size += 1;
551575
declared_vars.insert(var.clone());
552576
instructions.push(IntermediateInstruction::CounterHint {
553577
res_offset: compiler
@@ -590,7 +614,6 @@ fn handle_const_malloc(
590614
size: usize,
591615
label: &ConstMallocLabel,
592616
) {
593-
todo!(); // amend stack frame layout
594617
declared_vars.insert(var.clone());
595618
instructions.push(IntermediateInstruction::Computation {
596619
operation: Operation::Add,
@@ -600,7 +623,8 @@ fn handle_const_malloc(
600623
offset: compiler.get_offset(&var.clone().into()),
601624
},
602625
});
603-
compiler.const_mallocs.insert(*label, compiler.stack_size);
626+
let current_scope = compiler.stack_frame_layout.scopes.last_mut().unwrap();
627+
current_scope.const_mallocs.insert(*label, compiler.stack_size);
604628
compiler.stack_size += size;
605629
}
606630

@@ -693,6 +717,9 @@ fn find_internal_vars(lines: &[SimpleLine]) -> BTreeSet<Var> {
693717
let mut internal_vars = BTreeSet::new();
694718
for line in lines {
695719
match line {
720+
SimpleLine::ForwardDeclaration { var } => {
721+
internal_vars.insert(var.clone());
722+
}
696723
SimpleLine::Match { arms, .. } => {
697724
for arm in arms {
698725
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>,

‎crates/lean_compiler/src/lib.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ pub fn compile_program(program: String) -> Bytecode {
1616
let (parsed_program, function_locations) = parse_program(&program).unwrap();
1717
// println!("Parsed program: {}", parsed_program.to_string());
1818
let simple_program = simplify_program(parsed_program);
19-
// println!("Simplified program: {}", simple_program.to_string());
19+
println!("Simplified program: {}", simple_program.to_string());
2020
let intermediate_bytecode = compile_to_intermediate_bytecode(simple_program).unwrap();
2121
// println!("Intermediate Bytecode:\n\n{}", intermediate_bytecode.to_string());
2222

0 commit comments

Comments
 (0)