Skip to content

Commit 3785677

Browse files
TomWambsgansTomWambsgans
committed
wip
1 parent 3fcb104 commit 3785677

File tree

7 files changed

+149
-141
lines changed

7 files changed

+149
-141
lines changed

‎crates/lean_compiler/src/a_simplify_lang.rs‎

Lines changed: 82 additions & 86 deletions
Large diffs are not rendered by default.

‎crates/lean_compiler/src/b_compile_intermediate.rs‎

Lines changed: 19 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -63,33 +63,37 @@ impl Compiler {
6363
impl SimpleExpr {
6464
fn to_mem_after_fp_or_constant(&self, compiler: &Compiler) -> IntermediateValue {
6565
match self {
66-
Self::Var(var) => IntermediateValue::MemoryAfterFp {
66+
Self::Memory(VarOrConstMallocAccess::Var(var)) => IntermediateValue::MemoryAfterFp {
6767
offset: compiler.get_offset(&var.clone().into()),
6868
},
69+
Self::Memory(VarOrConstMallocAccess::ConstMallocAccess { malloc_label, offset }) => {
70+
IntermediateValue::MemoryAfterFp {
71+
offset: compiler.get_offset(&VarOrConstMallocAccess::ConstMallocAccess {
72+
malloc_label: *malloc_label,
73+
offset: offset.clone(),
74+
}),
75+
}
76+
}
6977
Self::Constant(c) => IntermediateValue::Constant(c.clone()),
70-
Self::ConstMallocAccess { malloc_label, offset } => IntermediateValue::MemoryAfterFp {
71-
offset: compiler.get_offset(&VarOrConstMallocAccess::ConstMallocAccess {
72-
malloc_label: *malloc_label,
73-
offset: offset.clone(),
74-
}),
75-
},
7678
}
7779
}
7880
}
7981

8082
impl IntermediateValue {
8183
fn from_simple_expr(expr: &SimpleExpr, compiler: &Compiler) -> Self {
8284
match expr {
83-
SimpleExpr::Var(var) => Self::MemoryAfterFp {
85+
SimpleExpr::Memory(VarOrConstMallocAccess::Var(var)) => Self::MemoryAfterFp {
8486
offset: compiler.get_offset(&var.clone().into()),
8587
},
88+
SimpleExpr::Memory(VarOrConstMallocAccess::ConstMallocAccess { malloc_label, offset }) => {
89+
Self::MemoryAfterFp {
90+
offset: compiler.get_offset(&VarOrConstMallocAccess::ConstMallocAccess {
91+
malloc_label: *malloc_label,
92+
offset: offset.clone(),
93+
}),
94+
}
95+
}
8696
SimpleExpr::Constant(c) => Self::Constant(c.clone()),
87-
SimpleExpr::ConstMallocAccess { malloc_label, offset } => Self::MemoryAfterFp {
88-
offset: compiler.get_offset(&VarOrConstMallocAccess::ConstMallocAccess {
89-
malloc_label: *malloc_label,
90-
offset: offset.clone(),
91-
}),
92-
},
9397
}
9498
}
9599

@@ -373,7 +377,7 @@ fn compile_lines(
373377
}
374378

375379
SimpleLine::RawAccess { res, index, shift } => {
376-
if let SimpleExpr::Var(var) = res
380+
if let SimpleExpr::Memory(VarOrConstMallocAccess::Var(var)) = res
377381
&& !compiler.is_in_scope(var)
378382
{
379383
let current_scope_layout = compiler.stack_frame_layout.scopes.last_mut().unwrap();

‎crates/lean_compiler/src/lang.rs‎

Lines changed: 9 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ use std::collections::{BTreeMap, BTreeSet};
55
use std::fmt::{Display, Formatter};
66
use utils::ToUsize;
77

8+
use crate::a_simplify_lang::VarOrConstMallocAccess;
89
use crate::{F, parser::ConstArrayValue};
910
pub use lean_vm::{FileId, FunctionName, SourceLocation};
1011

@@ -39,12 +40,8 @@ pub type ConstMallocLabel = usize;
3940

4041
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
4142
pub enum SimpleExpr {
42-
Var(Var),
43+
Memory(VarOrConstMallocAccess),
4344
Constant(ConstExpression),
44-
ConstMallocAccess {
45-
malloc_label: ConstMallocLabel,
46-
offset: ConstExpression,
47-
},
4845
}
4946

5047
impl SimpleExpr {
@@ -63,13 +60,6 @@ impl SimpleExpr {
6360
pub const fn is_constant(&self) -> bool {
6461
matches!(self, Self::Constant(_))
6562
}
66-
67-
pub fn simplify_if_const(&self) -> Self {
68-
if let Self::Constant(constant) = self {
69-
return constant.clone().into();
70-
}
71-
self.clone()
72-
}
7363
}
7464

7565
impl From<ConstantValue> for SimpleExpr {
@@ -86,16 +76,15 @@ impl From<ConstExpression> for SimpleExpr {
8676

8777
impl From<Var> for SimpleExpr {
8878
fn from(var: Var) -> Self {
89-
Self::Var(var)
79+
VarOrConstMallocAccess::Var(var).into()
9080
}
9181
}
9282

9383
impl SimpleExpr {
9484
pub fn as_constant(&self) -> Option<ConstExpression> {
9585
match self {
96-
Self::Var(_) => None,
9786
Self::Constant(constant) => Some(constant.clone()),
98-
Self::ConstMallocAccess { .. } => None,
87+
Self::Memory(_) => None,
9988
}
10089
}
10190

@@ -226,7 +215,7 @@ impl Display for Condition {
226215
pub enum Expression {
227216
Value(SimpleExpr),
228217
ArrayAccess {
229-
array: SimpleExpr,
218+
array: Var,
230219
index: Vec<Self>, // multi-dimensional array access
231220
},
232221
MathExpr(MathOperation, Vec<Self>),
@@ -356,10 +345,7 @@ impl Expression {
356345
self.eval_with(
357346
&|value: &SimpleExpr| value.as_constant()?.naive_eval(),
358347
&|arr, indexes| {
359-
let SimpleExpr::Var(name) = arr else {
360-
return None;
361-
};
362-
let array = const_arrays.get(name)?;
348+
let array = const_arrays.get(arr)?;
363349
assert_eq!(indexes.len(), array.depth());
364350
let idx = indexes.iter().map(|e| e.to_usize()).collect::<Vec<_>>();
365351
array.navigate(&idx)?.as_scalar().map(F::from_usize)
@@ -370,7 +356,7 @@ impl Expression {
370356
pub fn eval_with<ValueFn, ArrayFn>(&self, value_fn: &ValueFn, array_fn: &ArrayFn) -> Option<F>
371357
where
372358
ValueFn: Fn(&SimpleExpr) -> Option<F>,
373-
ArrayFn: Fn(&SimpleExpr, Vec<F>) -> Option<F>,
359+
ArrayFn: Fn(&Var, Vec<F>) -> Option<F>,
374360
{
375361
match self {
376362
Self::Value(value) => value_fn(value),
@@ -405,7 +391,7 @@ impl Expression {
405391
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
406392
pub enum AssignmentTarget {
407393
Var(Var),
408-
ArrayAccess { array: SimpleExpr, index: Box<Expression> },
394+
ArrayAccess { array: Var, index: Box<Expression> },
409395
}
410396

411397
impl Display for AssignmentTarget {
@@ -713,11 +699,8 @@ impl Display for ConstantValue {
713699
impl Display for SimpleExpr {
714700
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
715701
match self {
716-
Self::Var(var) => write!(f, "{var}"),
717702
Self::Constant(constant) => write!(f, "{constant}"),
718-
Self::ConstMallocAccess { malloc_label, offset } => {
719-
write!(f, "malloc_access({malloc_label}, {offset})")
720-
}
703+
Self::Memory(var_or_const_malloc_access) => write!(f, "{var_or_const_malloc_access}"),
721704
}
722705
}
723706
}

‎crates/lean_compiler/src/parser/parsers/expression.rs‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -80,16 +80,13 @@ pub struct ArrayAccessParser;
8080
impl Parse<Expression> for ArrayAccessParser {
8181
fn parse(&self, pair: ParsePair<'_>, ctx: &mut ParseContext) -> ParseResult<Expression> {
8282
let mut inner = pair.into_inner();
83-
let array_name = next_inner_pair(&mut inner, "array name")?.as_str().to_string();
83+
let array = next_inner_pair(&mut inner, "array name")?.as_str().to_string();
8484

8585
let index: Vec<Expression> = inner
8686
.map(|idx_pair| ExpressionParser.parse(idx_pair, ctx))
8787
.collect::<Result<Vec<_>, _>>()?;
8888

89-
Ok(Expression::ArrayAccess {
90-
array: SimpleExpr::Var(array_name),
91-
index,
92-
})
89+
Ok(Expression::ArrayAccess { array, index })
9390
}
9491
}
9592

‎crates/lean_compiler/src/parser/parsers/function.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use super::statement::StatementParser;
33
use super::{Parse, ParseContext, next_inner_pair};
44
use crate::{
55
SourceLineNumber,
6-
lang::{AssignmentTarget, Expression, Function, Line, SimpleExpr, SourceLocation},
6+
lang::{AssignmentTarget, Expression, Function, Line, SourceLocation},
77
parser::{
88
error::{ParseResult, SemanticError},
99
grammar::{ParsePair, Rule},
@@ -129,7 +129,7 @@ impl Parse<AssignmentTarget> for AssignmentTargetParser {
129129
let array = next_inner_pair(&mut inner_pairs, "array name")?.as_str().to_string();
130130
let index = ExpressionParser.parse(next_inner_pair(&mut inner_pairs, "array index")?, ctx)?;
131131
Ok(AssignmentTarget::ArrayAccess {
132-
array: SimpleExpr::Var(array),
132+
array,
133133
index: Box::new(index),
134134
})
135135
}

‎crates/lean_compiler/src/parser/parsers/literal.rs‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
use super::expression::ExpressionParser;
22
use super::{ConstArrayValue, Parse, ParseContext, ParsedConstant, next_inner_pair};
3+
use crate::a_simplify_lang::VarOrConstMallocAccess;
34
use crate::{
45
F,
56
lang::{ConstExpression, ConstantValue, SimpleExpr},
@@ -90,16 +91,13 @@ pub fn evaluate_const_expr(expr: &crate::lang::Expression, ctx: &ParseContext) -
9091
expr.eval_with(
9192
&|simple_expr| match simple_expr {
9293
SimpleExpr::Constant(cst) => cst.naive_eval(),
93-
SimpleExpr::Var(var) => ctx.get_constant(var).map(F::from_usize),
94-
SimpleExpr::ConstMallocAccess { .. } => None,
94+
SimpleExpr::Memory(VarOrConstMallocAccess::Var(var)) => ctx.get_constant(var).map(F::from_usize),
95+
SimpleExpr::Memory(VarOrConstMallocAccess::ConstMallocAccess { .. }) => None,
9596
},
9697
&|arr, index| {
9798
// Support const array access in expressions
98-
let SimpleExpr::Var(name) = arr else {
99-
return None;
100-
};
10199
let idx = index.iter().map(|e| e.to_usize()).collect::<Vec<_>>();
102-
let array = ctx.get_const_array(name)?;
100+
let array = ctx.get_const_array(arr)?;
103101
array.navigate(&idx)?.as_scalar().map(F::from_usize)
104102
},
105103
)
@@ -161,7 +159,7 @@ impl VarOrConstantParser {
161159
}
162160
// Otherwise treat as variable reference
163161
else {
164-
Ok(SimpleExpr::Var(text.to_string()))
162+
Ok(VarOrConstMallocAccess::Var(text.to_string()).into())
165163
}
166164
}
167165
}

‎crates/lean_compiler/tests/test_compiler.rs‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -456,6 +456,36 @@ fn test_mini_program_4() {
456456
dbg!(&poseidon24_permute(public_input)[16..]);
457457
}
458458

459+
#[test]
460+
fn test_mini_program_5() {
461+
let program = r#"
462+
fn main() {
463+
arr = malloc(10);
464+
arr[6] = 42;
465+
arr[8] = 11;
466+
sum_1 = func_1(arr[6], arr[8]);
467+
assert sum_1 == 53;
468+
return;
469+
}
470+
471+
fn func_1(i, j) inline -> 1 {
472+
for k in 0..i {
473+
for u in 0..j {
474+
assert k + u != 1000000;
475+
}
476+
}
477+
return i + j;
478+
}
479+
480+
"#;
481+
compile_and_run(
482+
&ProgramSource::Raw(program.to_string()),
483+
(&[], &[]),
484+
DEFAULT_NO_VEC_RUNTIME_MEMORY,
485+
false,
486+
);
487+
}
488+
459489
#[test]
460490
fn test_inlined() {
461491
let program = r#"

0 commit comments

Comments
 (0)