Skip to content

Commit a7892ec

Browse files
TomWambsgansweijiekoh
TomWambsgans
and
weijiekoh
committed
Co-authored-by: Koh Wei Jie <[email protected]>
1 parent bb2fbae commit a7892ec

File tree

9 files changed

+242
-13
lines changed

9 files changed

+242
-13
lines changed

‎crates/lean_compiler/src/a_simplify_lang.rs‎

Lines changed: 56 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,11 @@ pub enum SimpleLine {
123123
location: SourceLocation,
124124
},
125125
DebugAssert(BooleanExpr<SimpleExpr>, SourceLocation),
126+
/// Range check: assert val <= bound
127+
RangeCheck {
128+
val: SimpleExpr,
129+
bound: SimpleExpr,
130+
},
126131
}
127132

128133
impl SimpleLine {
@@ -157,7 +162,8 @@ impl SimpleLine {
157162
| Self::HintMAlloc { .. }
158163
| Self::ConstMalloc { .. }
159164
| Self::LocationReport { .. }
160-
| Self::DebugAssert(..) => vec![],
165+
| Self::DebugAssert(..)
166+
| Self::RangeCheck { .. } => vec![],
161167
}
162168
}
163169

@@ -182,7 +188,8 @@ impl SimpleLine {
182188
| Self::HintMAlloc { .. }
183189
| Self::ConstMalloc { .. }
184190
| Self::LocationReport { .. }
185-
| Self::DebugAssert(..) => vec![],
191+
| Self::DebugAssert(..)
192+
| Self::RangeCheck { .. } => vec![],
186193
}
187194
}
188195
}
@@ -2257,6 +2264,49 @@ fn simplify_lines(
22572264
};
22582265
res.push(SimpleLine::equality(var, other));
22592266
}
2267+
Boolean::LessThan => {
2268+
// assert left < right is equivalent to assert left <= right - 1
2269+
let bound_minus_one = state.counters.aux_var();
2270+
res.push(SimpleLine::Assignment {
2271+
var: bound_minus_one.clone().into(),
2272+
operation: MathOperation::Sub,
2273+
arg0: right,
2274+
arg1: SimpleExpr::one(),
2275+
});
2276+
2277+
// We add a debug assert for sanity
2278+
res.push(SimpleLine::DebugAssert(
2279+
BooleanExpr {
2280+
kind: Boolean::LessOrEqual,
2281+
left: left.clone(),
2282+
right: bound_minus_one.clone().into(),
2283+
},
2284+
*location,
2285+
));
2286+
2287+
res.push(SimpleLine::RangeCheck {
2288+
val: left,
2289+
bound: bound_minus_one.into(),
2290+
});
2291+
}
2292+
Boolean::LessOrEqual => {
2293+
// Range check: assert left <= right
2294+
2295+
// we add a debug assert for sanity
2296+
res.push(SimpleLine::DebugAssert(
2297+
BooleanExpr {
2298+
kind: Boolean::LessOrEqual,
2299+
left: left.clone(),
2300+
right: right.clone(),
2301+
},
2302+
*location,
2303+
));
2304+
2305+
res.push(SimpleLine::RangeCheck {
2306+
val: left,
2307+
bound: right,
2308+
});
2309+
}
22602310
}
22612311
}
22622312
}
@@ -2273,6 +2323,7 @@ fn simplify_lines(
22732323
let (left, right, then_branch, else_branch) = match condition.kind {
22742324
Boolean::Equal => (&condition.left, &condition.right, else_branch, then_branch), // switched
22752325
Boolean::Different => (&condition.left, &condition.right, then_branch, else_branch),
2326+
Boolean::LessThan | Boolean::LessOrEqual => unreachable!(),
22762327
};
22772328

22782329
let left_simplified = simplify_expr(ctx, state, const_malloc, left, &mut res)?;
@@ -3564,6 +3615,9 @@ impl SimpleLine {
35643615
Self::DebugAssert(bool, _) => {
35653616
format!("debug_assert({bool})")
35663617
}
3618+
Self::RangeCheck { val, bound } => {
3619+
format!("range_check({val} <= {bound})")
3620+
}
35673621
};
35683622
format!("{spaces}{line_str}")
35693623
}

‎crates/lean_compiler/src/b_compile_intermediate.rs‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -562,6 +562,81 @@ fn compile_lines(
562562
};
563563
instructions.push(IntermediateInstruction::DebugAssert(boolean_simplified, *location));
564564
}
565+
SimpleLine::RangeCheck { val, bound } => {
566+
// Range check for val <= bound compiles to:
567+
// 1. DEREF: m[fp + aux1] = m[m[fp + val_offset]] - proves val < M
568+
// 2. ADD: m[fp + val_offset] + m[fp + aux2] = bound - computes complement
569+
// 3. DEREF: m[fp + aux3] = m[m[fp + aux2]] - proves complement < M
570+
//
571+
// DerefHint records constraints: memory[target] = memory[memory[src]]
572+
// These are resolved at end of execution in correct order.
573+
574+
// Get the offset of the value being range-checked
575+
let val_offset = match val {
576+
SimpleExpr::Memory(var_or_const) => compiler.get_offset(var_or_const),
577+
SimpleExpr::Constant(val_const) => {
578+
// For constants, we need to store in a temp variable first
579+
let temp_offset = compiler.stack_pos;
580+
compiler.stack_pos += 1;
581+
instructions.push(IntermediateInstruction::Computation {
582+
operation: Operation::Add,
583+
arg_a: IntermediateValue::Constant(val_const.clone()),
584+
arg_b: IntermediateValue::Constant(ConstExpression::zero()),
585+
res: IntermediateValue::MemoryAfterFp {
586+
offset: ConstExpression::from_usize(temp_offset),
587+
},
588+
});
589+
ConstExpression::from_usize(temp_offset)
590+
}
591+
};
592+
593+
// Allocate 3 auxiliary cells
594+
let aux1_offset = ConstExpression::from_usize(compiler.stack_pos);
595+
compiler.stack_pos += 1;
596+
let aux2_offset = ConstExpression::from_usize(compiler.stack_pos);
597+
compiler.stack_pos += 1;
598+
let aux3_offset = ConstExpression::from_usize(compiler.stack_pos);
599+
compiler.stack_pos += 1;
600+
601+
// DerefHint for first DEREF: memory[aux1] = memory[memory[val_offset]]
602+
instructions.push(IntermediateInstruction::DerefHint {
603+
offset_src: val_offset.clone(),
604+
offset_target: aux1_offset.clone(),
605+
});
606+
607+
// 1. DEREF: m[fp + aux1] = m[m[fp + val_offset]]
608+
instructions.push(IntermediateInstruction::Deref {
609+
shift_0: val_offset.clone(),
610+
shift_1: ConstExpression::zero(),
611+
res: IntermediateValue::MemoryAfterFp { offset: aux1_offset },
612+
});
613+
614+
// 2. ADD: m[fp + val_offset] + m[fp + aux2] = bound
615+
let bound_value = IntermediateValue::from_simple_expr(bound, compiler);
616+
instructions.push(IntermediateInstruction::Computation {
617+
operation: Operation::Add,
618+
arg_a: IntermediateValue::MemoryAfterFp {
619+
offset: val_offset.clone(),
620+
},
621+
arg_b: IntermediateValue::MemoryAfterFp {
622+
offset: aux2_offset.clone(),
623+
},
624+
res: bound_value,
625+
});
626+
627+
// DerefHint for second DEREF: memory[aux3] = memory[memory[aux2]]
628+
instructions.push(IntermediateInstruction::DerefHint {
629+
offset_src: aux2_offset.clone(),
630+
offset_target: aux3_offset.clone(),
631+
});
632+
633+
// 3. DEREF: m[fp + aux3] = m[m[fp + aux2]]
634+
instructions.push(IntermediateInstruction::Deref {
635+
shift_0: aux2_offset,
636+
shift_1: ConstExpression::zero(),
637+
res: IntermediateValue::MemoryAfterFp { offset: aux3_offset },
638+
});
639+
}
565640
}
566641
}
567642

‎crates/lean_compiler/src/c_compile_final.rs‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ impl IntermediateInstruction {
1313
| Self::Inverse { .. }
1414
| Self::LocationReport { .. }
1515
| Self::DebugAssert { .. }
16+
| Self::DerefHint { .. }
1617
| Self::PanicHint { .. } => true,
1718
Self::Computation { .. }
1819
| Self::Panic
@@ -349,6 +350,16 @@ fn compile_block(
349350
);
350351
hints.entry(pc).or_default().push(hint);
351352
}
353+
IntermediateInstruction::DerefHint {
354+
offset_src,
355+
offset_target,
356+
} => {
357+
let hint = Hint::DerefHint {
358+
offset_src: eval_const_expression_usize(&offset_src, compiler),
359+
offset_target: eval_const_expression_usize(&offset_target, compiler),
360+
};
361+
hints.entry(pc).or_default().push(hint);
362+
}
352363
IntermediateInstruction::PanicHint { message } => {
353364
let hint = Hint::Panic { message };
354365
hints.entry(pc).or_default().push(hint);

‎crates/lean_compiler/src/ir/instruction.rs‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,13 @@ pub enum IntermediateInstruction {
4646
size: IntermediateValue, // the hint
4747
},
4848
CustomHint(CustomHint, Vec<IntermediateValue>),
49+
/// Deref hint for range checks - records constraint resolved at end of execution
50+
DerefHint {
51+
/// Offset of cell containing the address to dereference
52+
offset_src: ConstExpression,
53+
/// Offset of cell where result will be stored
54+
offset_target: ConstExpression,
55+
},
4956
Print {
5057
line_info: String, // information about the line where the print occurs
5158
content: Vec<IntermediateValue>, // values to print
@@ -183,6 +190,12 @@ impl Display for IntermediateInstruction {
183190
Self::DebugAssert(boolean_expr, _) => {
184191
write!(f, "debug_assert {boolean_expr}")
185192
}
193+
Self::DerefHint {
194+
offset_src,
195+
offset_target,
196+
} => {
197+
write!(f, "m[fp + {offset_target}] = m[m[fp + {offset_src}]]")
198+
}
186199
Self::PanicHint { message } => match message {
187200
Some(msg) => write!(f, "panic hint: \"{msg}\""),
188201
None => write!(f, "panic hint"),

‎crates/lean_compiler/src/lang.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -277,6 +277,8 @@ impl Condition {
277277
Some(match cmp.kind {
278278
Boolean::Equal => left == right,
279279
Boolean::Different => left != right,
280+
Boolean::LessThan => left < right,
281+
Boolean::LessOrEqual => left <= right,
280282
})
281283
}
282284
}

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,8 @@ impl ComparisonParser {
175175
let kind = match op.as_str() {
176176
"==" => Boolean::Equal,
177177
"!=" => Boolean::Different,
178+
"<" => Boolean::LessThan,
179+
"<=" => Boolean::LessOrEqual,
178180
_ => unreachable!(),
179181
};
180182

‎crates/lean_vm/src/diagnostics/profiler.rs‎

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -66,17 +66,9 @@ pub(crate) fn profiling_report(
6666
report.push_str("──────────────────────────────────────────────────────────────────────────\n");
6767

6868
for (func_name, stats) in &function_data {
69-
let avg_exclusive = if stats.call_count > 0 {
70-
stats.exclusive_cycles / stats.call_count
71-
} else {
72-
0
73-
};
69+
let avg_exclusive = stats.exclusive_cycles.checked_div(stats.call_count).unwrap_or(0);
7470

75-
let avg_inclusive = if stats.call_count > 0 {
76-
stats.inclusive_cycles / stats.call_count
77-
} else {
78-
0
79-
};
71+
let avg_inclusive = stats.inclusive_cycles.checked_div(stats.call_count).unwrap_or(0);
8072

8173
report.push_str(&format!(
8274
"{:>5} │ {:>8} ┊ {:>8} │ {:>8} ┊ {:>8} │ {}\n",

‎crates/lean_vm/src/execution/runner.rs‎

Lines changed: 49 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ use crate::{
1313
};
1414
use multilinear_toolkit::prelude::*;
1515
use std::collections::{BTreeMap, BTreeSet};
16-
use utils::{poseidon16_permute, pretty_integer};
16+
use utils::{ToUsize, poseidon16_permute, pretty_integer};
1717
use xmss::Poseidon16History;
1818

1919
/// Number of instructions to show in stack trace
@@ -151,6 +151,45 @@ fn print_instruction_cycle_counts(bytecode: &Bytecode, pcs: Vec<CodeAddress>) {
151151
println!();
152152
}
153153

154+
/// Resolve pending deref hints in correct order
155+
///
156+
/// Each constraint has form: memory[target_addr] = memory[memory[src_addr]]
157+
/// Order matters because some src addresses might point to targets of other hints.
158+
/// We iteratively resolve constraints until no more progress, then fill remaining with 0.
159+
fn resolve_deref_hints(memory: &mut Memory, pending: &[(usize, usize)]) {
160+
let mut resolved: BTreeSet<usize> = BTreeSet::new();
161+
162+
loop {
163+
let mut made_progress = false;
164+
165+
for &(target_addr, src_addr) in pending {
166+
if resolved.contains(&target_addr) {
167+
continue;
168+
}
169+
let Some(addr) = memory.0.get(src_addr).copied().flatten() else {
170+
continue;
171+
};
172+
let Some(value) = memory.0.get(addr.to_usize()).copied().flatten() else {
173+
continue;
174+
};
175+
memory.set(target_addr, value).unwrap();
176+
resolved.insert(target_addr);
177+
made_progress = true;
178+
}
179+
180+
if !made_progress {
181+
break;
182+
}
183+
}
184+
185+
// Fill any remaining unresolved targets with 0
186+
for &(target_addr, _src_addr) in pending {
187+
if !resolved.contains(&target_addr) {
188+
let _ = memory.set(target_addr, F::ZERO);
189+
}
190+
}
191+
}
192+
154193
/// Helper function that performs the actual bytecode execution
155194
#[allow(clippy::too_many_arguments)] // TODO
156195
fn execute_bytecode_helper(
@@ -206,6 +245,9 @@ fn execute_bytecode_helper(
206245
let mut counter_hint = 0;
207246
let mut cpu_cycles_before_new_line = 0;
208247

248+
// Pending deref hints: (target_addr, src_addr) constraints to resolve at end
249+
let mut pending_deref_hints: Vec<(usize, usize)> = Vec::new();
250+
209251
while pc != ENDING_PC {
210252
if pc >= bytecode.instructions.len() {
211253
return Err((pc, RunnerError::PCOutOfBounds));
@@ -231,6 +273,7 @@ fn execute_bytecode_helper(
231273
checkpoint_ap: &mut checkpoint_ap,
232274
profiling,
233275
memory_profile: &mut mem_profile,
276+
pending_deref_hints: &mut pending_deref_hints,
234277
};
235278
hint.execute_hint(&mut hint_ctx).map_err(|e| (pc, e))?;
236279
}
@@ -254,6 +297,11 @@ fn execute_bytecode_helper(
254297
.map_err(|e| (pc, e))?;
255298
}
256299

300+
// Resolve pending deref hints in correct order
301+
// Constraint: memory[target_addr] = memory[memory[src_addr]]
302+
// Order matters because some src addresses might point to targets of other hints
303+
resolve_deref_hints(&mut memory, &pending_deref_hints);
304+
257305
assert_eq!(
258306
n_poseidon16_precomputed_used,
259307
poseidons_16_precomputed.len(),

0 commit comments

Comments
 (0)