Skip to content

Commit c6f9fd4

Browse files
TomWambsgansTomWambsgans
committed
zkDSL: additional optimizations after a_simplify_lang 
1 parent 6fdcb8f commit c6f9fd4

3 files changed

Lines changed: 156 additions & 111 deletions

File tree

‎crates/lean_compiler/src/a_simplify_lang/mod.rs‎

Lines changed: 28 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use crate::{
22
CompilationFlags, F,
3-
a_simplify_lang::copy_propagation::propagate_copies,
3+
a_simplify_lang::post_optimization::propagate_copies,
44
lang::*,
55
parser::{ConstArrayValue, parse_program},
66
};
@@ -15,7 +15,7 @@ use std::{
1515
};
1616
use utils::{Counter, ToUsize};
1717

18-
mod copy_propagation;
18+
mod post_optimization;
1919

2020
#[derive(Debug, Clone)]
2121
pub struct SimpleProgram {
@@ -75,8 +75,9 @@ pub enum SimpleLine {
7575
var: Var,
7676
},
7777
Assignment {
78-
var: VarOrConstMallocAccess,
79-
operation: MathOperation, // add / sub / div / mul
78+
// `var = arg0 op arg1`
79+
var: SimpleExpr,
80+
op: MathOperation, // add / sub / div / mul
8081
arg0: SimpleExpr,
8182
arg1: SimpleExpr,
8283
},
@@ -150,10 +151,10 @@ pub enum SimpleLine {
150151
}
151152

152153
impl SimpleLine {
153-
pub fn equality(arg0: impl Into<VarOrConstMallocAccess>, arg1: impl Into<SimpleExpr>) -> Self {
154+
pub fn equality(arg0: impl Into<SimpleExpr>, arg1: impl Into<SimpleExpr>) -> Self {
154155
SimpleLine::Assignment {
155156
var: arg0.into(),
156-
operation: MathOperation::Add,
157+
op: MathOperation::Add,
157158
arg0: arg1.into(),
158159
arg1: SimpleExpr::zero(),
159160
}
@@ -2456,7 +2457,7 @@ fn simplify_lines(
24562457
}
24572458
res.push(SimpleLine::Assignment {
24582459
var: target_var.into(),
2459-
operation: *operation,
2460+
op: *operation,
24602461
arg0: args_simplified[0].clone(),
24612462
arg1: args_simplified[1].clone(),
24622463
});
@@ -2482,7 +2483,7 @@ fn simplify_lines(
24822483
if let SimpleExpr::Constant(offset) = &simplified_index
24832484
&& let Some(array_name) = array.as_var()
24842485
&& let Some(label) = const_malloc.map.get(array_name)
2485-
&& let Expression::MathExpr(operation, args) = value
2486+
&& let Expression::MathExpr(op, args) = value
24862487
{
24872488
let var = VarOrConstMallocAccess::ConstMallocAccess {
24882489
malloc_label: *label,
@@ -2494,18 +2495,18 @@ fn simplify_lines(
24942495
.collect::<Result<Vec<_>, _>>()?;
24952496
// If all operands are constants, evaluate at compile time
24962497
if let Some(const_args) = SimpleExpr::try_vec_as_constant(&simplified_args) {
2497-
let result = ConstExpression::MathExpr(*operation, const_args);
2498+
let result = ConstExpression::MathExpr(*op, const_args);
24982499
res.push(SimpleLine::equality(var, SimpleExpr::Constant(result)));
24992500
} else {
2500-
if !operation.supports_runtime() {
2501+
if !op.supports_runtime() {
25012502
return Err(format!(
2502-
"Operation `{operation}` is compile-time only; all operands must be constants"
2503+
"Operation `{op}` is compile-time only; all operands must be constants"
25032504
));
25042505
}
25052506
assert_eq!(simplified_args.len(), 2);
25062507
res.push(SimpleLine::Assignment {
2507-
var,
2508-
operation: *operation,
2508+
var: var.into(),
2509+
op: *op,
25092510
arg0: simplified_args[0].clone(),
25102511
arg1: simplified_args[1].clone(),
25112512
});
@@ -2549,7 +2550,7 @@ fn simplify_lines(
25492550
let diff_var = state.counters.aux_var();
25502551
res.push(SimpleLine::Assignment {
25512552
var: diff_var.clone().into(),
2552-
operation: MathOperation::Sub,
2553+
op: MathOperation::Sub,
25532554
arg0: left,
25542555
arg1: right,
25552556
});
@@ -2592,7 +2593,7 @@ fn simplify_lines(
25922593
let bound_minus_one = state.counters.aux_var();
25932594
res.push(SimpleLine::Assignment {
25942595
var: bound_minus_one.clone().into(),
2595-
operation: MathOperation::Sub,
2596+
op: MathOperation::Sub,
25962597
arg0: right,
25972598
arg1: SimpleExpr::one(),
25982599
});
@@ -2652,7 +2653,7 @@ fn simplify_lines(
26522653
let diff_var = state.counters.aux_var();
26532654
res.push(SimpleLine::Assignment {
26542655
var: diff_var.clone().into(),
2655-
operation: MathOperation::Sub,
2656+
op: MathOperation::Sub,
26562657
arg0: left_simplified,
26572658
arg1: right_simplified,
26582659
});
@@ -3014,24 +3015,24 @@ fn simplify_expr(
30143015
);
30153016
Ok(VarOrConstMallocAccess::Var(aux_arr).into())
30163017
}
3017-
Expression::MathExpr(operation, args) => {
3018+
Expression::MathExpr(op, args) => {
30183019
let simplified_args = args
30193020
.iter()
30203021
.map(|arg| simplify_expr(ctx, state, const_malloc, arg, lines))
30213022
.collect::<Result<Vec<_>, _>>()?;
30223023
if let Some(const_args) = SimpleExpr::try_vec_as_constant(&simplified_args) {
3023-
return Ok(SimpleExpr::Constant(ConstExpression::MathExpr(*operation, const_args)));
3024+
return Ok(SimpleExpr::Constant(ConstExpression::MathExpr(*op, const_args)));
30243025
}
3025-
if !operation.supports_runtime() {
3026+
if !op.supports_runtime() {
30263027
return Err(format!(
3027-
"Operation `{operation}` is compile-time only; all operands must be constants"
3028+
"Operation `{op}` is compile-time only; all operands must be constants"
30283029
));
30293030
}
30303031
let aux_var = state.counters.aux_var();
30313032
assert_eq!(simplified_args.len(), 2);
30323033
lines.push(SimpleLine::Assignment {
30333034
var: aux_var.clone().into(),
3034-
operation: *operation,
3035+
op: *op,
30353036
arg0: simplified_args[0].clone(),
30363037
arg1: simplified_args[1].clone(),
30373038
});
@@ -3527,7 +3528,7 @@ fn handle_array_assignment(
35273528
let base_var = state.counters.aux_var();
35283529
res.push(SimpleLine::Assignment {
35293530
var: base_var.clone().into(),
3530-
operation: MathOperation::Add,
3531+
op: MathOperation::Add,
35313532
arg0: array.clone(),
35323533
arg1: SimpleExpr::zero(),
35333534
});
@@ -3549,7 +3550,7 @@ fn handle_array_assignment(
35493550
let ptr_var = state.counters.aux_var();
35503551
res.push(SimpleLine::Assignment {
35513552
var: ptr_var.clone().into(),
3552-
operation: MathOperation::Add,
3553+
op: MathOperation::Add,
35533554
arg0: base_addr,
35543555
arg1: simplified_index,
35553556
});
@@ -3580,7 +3581,7 @@ fn create_recursive_function(
35803581
let next_iter = format!("@incremented_{iterator}");
35813582
body.push(SimpleLine::Assignment {
35823583
var: next_iter.clone().into(),
3583-
operation: MathOperation::Add,
3584+
op: MathOperation::Add,
35843585
arg0: iterator.clone().into(),
35853586
arg1: SimpleExpr::one(),
35863587
});
@@ -3602,7 +3603,7 @@ fn create_recursive_function(
36023603
let instructions = vec![
36033604
SimpleLine::Assignment {
36043605
var: diff_var.clone().into(),
3605-
operation: MathOperation::Sub,
3606+
op: MathOperation::Sub,
36063607
arg0: iterator.into(),
36073608
arg1: end,
36083609
},
@@ -3960,13 +3961,8 @@ impl SimpleLine {
39603961
format!("match {value} {{\n{arms_str}\n{spaces}}}")
39613962
}
39623963

3963-
Self::Assignment {
3964-
var,
3965-
operation,
3966-
arg0,
3967-
arg1,
3968-
} => {
3969-
format!("{var} = {arg0} {operation} {arg1}")
3964+
Self::Assignment { var, op, arg0, arg1 } => {
3965+
format!("{var} = {arg0} {op} {arg1}")
39703966
}
39713967
Self::CustomHint(hint, args) => {
39723968
format!(

0 commit comments

Comments
 (0)