From d0c2b0e223eab4379800091a238cd6042c036da4 Mon Sep 17 00:00:00 2001 From: Koh Wei Jie Date: Tue, 28 Oct 2025 20:17:45 -0700 Subject: [PATCH 01/10] range check implemented in compiler; tests WIP --- Cargo.lock | 1 + crates/lean_compiler/src/a_simplify_lang.rs | 26 ++- .../src/b_compile_intermediate.rs | 69 +++++++- crates/lean_compiler/src/c_compile_final.rs | 16 +- crates/lean_compiler/src/grammar.pest | 7 +- crates/lean_compiler/src/ir/instruction.rs | 11 +- crates/lean_compiler/src/lang.rs | 5 + .../src/parser/parsers/statement.rs | 22 ++- crates/lean_compiler/tests/test_compiler.rs | 39 +++++ crates/lean_prover/Cargo.toml | 3 +- crates/lean_prover/tests/test_range_check.rs | 159 ++++++++++++++++++ .../witness_generation/src/execution_trace.rs | 1 + .../src/instruction_encoder.rs | 1 + crates/lean_vm/src/isa/hint.rs | 12 +- crates/lean_vm/src/isa/instruction.rs | 19 ++- 15 files changed, 372 insertions(+), 19 deletions(-) create mode 100644 crates/lean_prover/tests/test_range_check.rs diff --git a/Cargo.lock b/Cargo.lock index f734ff0d..bed5b709 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -403,6 +403,7 @@ dependencies = [ "pest_derive", "poseidon_circuit", "rand", + "rand_chacha", "rayon", "tracing", "utils", diff --git a/crates/lean_compiler/src/a_simplify_lang.rs b/crates/lean_compiler/src/a_simplify_lang.rs index 558ff452..e7fa0be0 100644 --- a/crates/lean_compiler/src/a_simplify_lang.rs +++ b/crates/lean_compiler/src/a_simplify_lang.rs @@ -154,6 +154,10 @@ pub enum SimpleLine { LocationReport { location: SourceLineNumber, }, + RangeCheck { + value: SimpleExpr, + max: ConstExpression, + }, } pub fn simplify_program(mut program: Program) -> SimpleProgram { @@ -738,6 +742,14 @@ fn simplify_lines( location: *location, }); } + Line::RangeCheck { value, max } => { + let simplified_value = + simplify_expr(value, &mut res, counters, array_manager, const_malloc); + res.push(SimpleLine::RangeCheck { + value: simplified_value, + max: max.clone(), + }); + } } } @@ -958,7 +970,7 @@ pub fn find_variable_usage(lines: &[Line]) -> (BTreeSet, BTreeSet) { on_new_expr(index, &internal_vars, &mut external_vars); on_new_expr(value, &internal_vars, &mut external_vars); } - Line::Panic | Line::Break | Line::LocationReport { .. } => {} + Line::Panic | Line::Break | Line::LocationReport { .. } | Line::RangeCheck { .. } => {} } } @@ -1127,7 +1139,7 @@ pub fn inline_lines( inline_expr(index, args, inlining_count); inline_expr(value, args, inlining_count); } - Line::Panic | Line::Break | Line::LocationReport { .. } => {} + Line::Panic | Line::Break | Line::LocationReport { .. } | Line::RangeCheck { .. } => {} } } for (i, new_lines) in lines_to_replace.into_iter().rev() { @@ -1612,7 +1624,7 @@ fn replace_vars_for_unroll( Line::CounterHint { var } => { *var = format!("@unrolled_{unroll_index}_{iterator_value}_{var}"); } - Line::Break | Line::Panic | Line::LocationReport { .. } => {} + Line::Break | Line::Panic | Line::LocationReport { .. } | Line::RangeCheck { .. } => {} } } } @@ -2002,7 +2014,8 @@ fn get_function_called(lines: &[Line], function_called: &mut Vec) { | Line::MAlloc { .. } | Line::Panic | Line::Break - | Line::LocationReport { .. } => {} + | Line::LocationReport { .. } + | Line::RangeCheck { .. } => {} } } } @@ -2110,7 +2123,7 @@ fn replace_vars_by_const_in_lines(lines: &mut [Line], map: &BTreeMap) { assert!(!map.contains_key(var), "Variable {var} is a constant"); replace_vars_by_const_in_expr(size, map); } - Line::Panic | Line::Break | Line::LocationReport { .. } => {} + Line::Panic | Line::Break | Line::LocationReport { .. } | Line::RangeCheck { .. } => {} } } } @@ -2305,6 +2318,9 @@ impl SimpleLine { } Self::Panic => "panic".to_string(), Self::LocationReport { .. } => Default::default(), + Self::RangeCheck { value, max } => { + format!("range_check({value}, {max})") + } }; format!("{spaces}{line_str}") } diff --git a/crates/lean_compiler/src/b_compile_intermediate.rs b/crates/lean_compiler/src/b_compile_intermediate.rs index 92818757..899eadf4 100644 --- a/crates/lean_compiler/src/b_compile_intermediate.rs +++ b/crates/lean_compiler/src/b_compile_intermediate.rs @@ -1,6 +1,6 @@ use crate::{F, a_simplify_lang::*, ir::*, lang::*, precompiles::*}; use lean_vm::*; -use p3_field::Field; +use p3_field::{Field, PrimeCharacteristicRing}; use std::{ borrow::Borrow, collections::{BTreeMap, BTreeSet}, @@ -395,6 +395,7 @@ fn compile_lines( shift_0, shift_1: shift.clone(), res: res.to_mem_after_fp_or_constant(compiler), + for_range_check: false, }); } @@ -432,6 +433,7 @@ fn compile_lines( res: IntermediaryMemOrFpOrConstant::MemoryAfterFp { offset: compiler.get_offset(&ret_var.clone().into()), }, + for_range_check: false, }); } @@ -612,6 +614,65 @@ fn compile_lines( location: *location, }); } + SimpleLine::RangeCheck { value, max } => { + let x = match IntermediateValue::from_simple_expr(value, compiler) { + IntermediateValue::MemoryAfterFp { offset } => offset.naive_eval().unwrap(), + value::IntermediateValue::Fp => F::ZERO, + value::IntermediateValue::Constant(_) => unimplemented!(), + }; + + let t = max.naive_eval().unwrap(); + let aux_i: usize = compiler.stack_size + 0; + let aux_j: usize = compiler.stack_size + 1; + let aux_k: usize = compiler.stack_size + 2; + + // Step 1: DEREF: m[fp + i] == m[m[fp + x]] + // DEREF: m[fp + i] == m[value] + + let step_1 = IntermediateInstruction::Deref { + shift_0: ConstExpression::scalar(x.to_usize()), + shift_1: ConstExpression::from(0), + res: IntermediaryMemOrFpOrConstant::MemoryAfterFp { offset: aux_i.into() }, + for_range_check: true, + }; + + + // Step 2: ADD: m[m[fp + x]] + m[fp + j] == (t-1) + // m[fp + j] == t - 1 - value + // + // m[fp + j] == t - 1 - m[fp + x] + let q = t - F::ONE; + let step_2 = IntermediateInstruction::Computation { + operation: Operation::Add, + arg_a: IntermediateValue::MemoryAfterFp { offset: x.to_usize().into() }, + arg_c: IntermediateValue::MemoryAfterFp { offset: aux_j.into() }, // solve + res: IntermediateValue::Constant(q.to_usize().into()), // t - 1 + }; + + // Step 3: DEREF: m[fp + k] == m[m[fp + j]] + let step_3 = IntermediateInstruction::Deref { + shift_0: ConstExpression::scalar(aux_j), + shift_1: ConstExpression::from(0), + res: IntermediaryMemOrFpOrConstant::MemoryAfterFp { offset: aux_k.into() }, + for_range_check: true, + }; + + + // TODO: handle undefined memory access error + + //println!("aux_i: {}; {:?}", aux_i, step_1); + //println!("aux_j: {}; {:?}", aux_j, step_2); + //println!("aux_k: {}; {:?}", aux_k, step_3); + + instructions.push(IntermediateInstruction::RangeCheck { + value: IntermediateValue::from_simple_expr(value, compiler), + max: max.clone(), + }); + instructions.push(step_1); + instructions.push(step_2); + instructions.push(step_3); + compiler.stack_size += 3; + } } } @@ -689,11 +750,13 @@ fn setup_function_call( res: IntermediaryMemOrFpOrConstant::Constant(ConstExpression::label( return_label.clone(), )), + for_range_check: false, }, IntermediateInstruction::Deref { shift_0: new_fp_pos.into(), shift_1: ConstExpression::one(), res: IntermediaryMemOrFpOrConstant::Fp, + for_range_check: false, }, ]; @@ -703,6 +766,7 @@ fn setup_function_call( shift_0: new_fp_pos.into(), shift_1: (2 + i).into(), res: arg.to_mem_after_fp_or_constant(compiler), + for_range_check: false, }); } @@ -811,7 +875,8 @@ fn find_internal_vars(lines: &[SimpleLine]) -> BTreeSet { | SimpleLine::Print { .. } | SimpleLine::FunctionRet { .. } | SimpleLine::Precompile { .. } - | SimpleLine::LocationReport { .. } => {} + | SimpleLine::LocationReport { .. } + | SimpleLine::RangeCheck { .. } => {} } } internal_vars diff --git a/crates/lean_compiler/src/c_compile_final.rs b/crates/lean_compiler/src/c_compile_final.rs index a4fb8800..30d11f1f 100644 --- a/crates/lean_compiler/src/c_compile_final.rs +++ b/crates/lean_compiler/src/c_compile_final.rs @@ -13,7 +13,8 @@ impl IntermediateInstruction { | Self::DecomposeCustom { .. } | Self::CounterHint { .. } | Self::Inverse { .. } - | Self::LocationReport { .. } => true, + | Self::LocationReport { .. } + | Self::RangeCheck { .. } => true, Self::Computation { .. } | Self::Panic | Self::Deref { .. } @@ -22,7 +23,7 @@ impl IntermediateInstruction { | Self::Poseidon2_16 { .. } | Self::Poseidon2_24 { .. } | Self::DotProduct { .. } - | Self::MultilinearEval { .. } => false, + | Self::MultilinearEval { .. } => false } } } @@ -238,6 +239,7 @@ fn compile_block( shift_0, shift_1, res, + for_range_check, } => { low_level_bytecode.push(Instruction::Deref { shift_0: eval_const_expression(&shift_0, compiler).to_usize(), @@ -253,6 +255,7 @@ fn compile_block( MemOrFpOrConstant::Constant(eval_const_expression(&c, compiler)) } }, + for_range_check, }); } IntermediateInstruction::JumpIfNotZero { @@ -380,8 +383,15 @@ fn compile_block( let hint = Hint::LocationReport { location }; hints.entry(pc).or_default().push(hint); } + IntermediateInstruction::RangeCheck { value, max } => { + let hint = Hint::RangeCheck { + value: value.try_into_mem_or_fp(compiler).unwrap(), + // TODO: support max being an IntermediateValue + max: MemOrConstant::Constant(eval_const_expression(&max, compiler)), + }; + hints.entry(pc).or_default().push(hint); + } } - if !instruction.is_hint() { pc += 1; } diff --git a/crates/lean_compiler/src/grammar.pest b/crates/lean_compiler/src/grammar.pest index d6fc7d1a..61282407 100644 --- a/crates/lean_compiler/src/grammar.pest +++ b/crates/lean_compiler/src/grammar.pest @@ -24,9 +24,10 @@ statement = { return_statement | break_statement | continue_statement | - function_call | + range_check_statement | assert_eq_statement | - assert_not_eq_statement + assert_not_eq_statement | + function_call // Placed at the end so that it doesn't override other statements like range_check and assert_eq } return_statement = { "return" ~ (tuple_expression)? ~ ";" } @@ -61,6 +62,8 @@ var_list = { identifier ~ ("," ~ identifier)* } assert_eq_statement = { "assert" ~ add_expr ~ "==" ~ add_expr ~ ";" } assert_not_eq_statement = { "assert" ~ add_expr ~ "!=" ~ add_expr ~ ";" } +range_check_statement = { "range_check" ~ "(" ~ expression ~ "," ~ expression ~ ")" ~ ";" } + // Expressions tuple_expression = { expression ~ ("," ~ expression)* } expression = { neq_expr } diff --git a/crates/lean_compiler/src/ir/instruction.rs b/crates/lean_compiler/src/ir/instruction.rs index 40e73770..d2ff001a 100644 --- a/crates/lean_compiler/src/ir/instruction.rs +++ b/crates/lean_compiler/src/ir/instruction.rs @@ -17,6 +17,7 @@ pub enum IntermediateInstruction { shift_0: ConstExpression, shift_1: ConstExpression, res: IntermediaryMemOrFpOrConstant, + for_range_check: bool, }, // res = m[m[fp + shift_0]] Panic, Jump { @@ -86,6 +87,10 @@ pub enum IntermediateInstruction { LocationReport { location: SourceLineNumber, }, + RangeCheck { + value: IntermediateValue, + max: ConstExpression, + }, } impl IntermediateInstruction { @@ -152,7 +157,8 @@ impl Display for IntermediateInstruction { shift_0, shift_1, res, - } => write!(f, "{res} = m[m[fp + {shift_0}] + {shift_1}]"), + for_range_check, + } => write!(f, "{res} = m[m[fp + {shift_0}] + {shift_1}] for_range_check: {for_range_check}"), Self::Panic => write!(f, "panic"), Self::Jump { dest, updated_fp } => { if let Some(fp) = updated_fp { @@ -256,6 +262,9 @@ impl Display for IntermediateInstruction { Ok(()) } Self::LocationReport { .. } => Ok(()), + Self::RangeCheck { value, max } => { + write!(f, "range_check({value}, {max})") + } } } } diff --git a/crates/lean_compiler/src/lang.rs b/crates/lean_compiler/src/lang.rs index 7eea1ea0..8d9234cf 100644 --- a/crates/lean_compiler/src/lang.rs +++ b/crates/lean_compiler/src/lang.rs @@ -390,6 +390,10 @@ pub enum Line { LocationReport { location: SourceLineNumber, }, + RangeCheck { + value: Expression, + max: ConstExpression, + }, } impl Display for Expression { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { @@ -585,6 +589,7 @@ impl Line { } Self::Break => "break".to_string(), Self::Panic => "panic".to_string(), + Self::RangeCheck { value, max } => format!("range_check({value}, {max})"), }; format!("{spaces}{line_str}") } diff --git a/crates/lean_compiler/src/parser/parsers/statement.rs b/crates/lean_compiler/src/parser/parsers/statement.rs index 2d4e0fa9..b8a44d4f 100644 --- a/crates/lean_compiler/src/parser/parsers/statement.rs +++ b/crates/lean_compiler/src/parser/parsers/statement.rs @@ -4,7 +4,7 @@ use super::literal::ConstExprParser; use super::{Parse, ParseContext, next_inner_pair}; use crate::{ ir::HighLevelOperation, - lang::{AssumeBoolean, Boolean, Condition, Expression, Line}, + lang::{AssumeBoolean, Boolean, Condition, ConstExpression, Expression, Line}, parser::{ error::{ParseResult, SemanticError}, grammar::{ParsePair, Rule}, @@ -28,6 +28,7 @@ impl Parse for StatementParser { Rule::function_call => FunctionCallParser::parse(inner, ctx), Rule::assert_eq_statement => AssertEqParser::parse(inner, ctx), Rule::assert_not_eq_statement => AssertNotEqParser::parse(inner, ctx), + Rule::range_check_statement => RangeCheckParser::parse(inner, ctx), Rule::break_statement => Ok(Line::Break), Rule::continue_statement => { Err(SemanticError::new("Continue statement not implemented yet").into()) @@ -318,3 +319,22 @@ impl Parse for AssertNotEqParser { Ok(Line::Assert(Boolean::Different { left, right })) } } + +/// Parser for range check statements. +pub struct RangeCheckParser; + +impl Parse for RangeCheckParser { + fn parse(pair: ParsePair<'_>, ctx: &mut ParseContext) -> ParseResult { + let mut inner = pair.into_inner(); + let value = + ExpressionParser::parse(next_inner_pair(&mut inner, "range check value")?, ctx)?; + let max_expr = + ExpressionParser::parse(next_inner_pair(&mut inner, "range check max")?, ctx)?; + + // Convert the max expression to a const expression + let max = ConstExpression::try_from(max_expr) + .map_err(|_| SemanticError::new("Range check maximum must be a constant expression"))?; + + Ok(Line::RangeCheck { value, max }) + } +} diff --git a/crates/lean_compiler/tests/test_compiler.rs b/crates/lean_compiler/tests/test_compiler.rs index 86e56670..bec66b9d 100644 --- a/crates/lean_compiler/tests/test_compiler.rs +++ b/crates/lean_compiler/tests/test_compiler.rs @@ -4,6 +4,45 @@ use utils::{poseidon16_permute, poseidon24_permute}; const DEFAULT_NO_VEC_RUNTIME_MEMORY: usize = 1 << 15; +// TODO: create more test programs +fn range_check_program(value: usize, max: usize) -> String { + let program = format!( + r#" + //fn func(val) {{ + //if 0 == 0 {{ + //range_check(val, {max}); + //}} + //abc = 0; + //range_check(abc, {max}); + //return; + //}} + + fn main() {{ + val = {value}; + //func(val); + range_check(val, {max}); + //range_check(val, {max}); + return; + }} + "# + ); + program.to_string() +} + +#[test] +fn test_compile_range_check() { + let program = range_check_program(1000, 100000); + let bytecode = compile_program(program.clone()); + println!("{}", bytecode); + compile_and_run( + program.to_string(), + (&[], &[]), + DEFAULT_NO_VEC_RUNTIME_MEMORY, + false, + ); + //TODO test more +} + #[test] #[should_panic] fn test_duplicate_function_name() { diff --git a/crates/lean_prover/Cargo.toml b/crates/lean_prover/Cargo.toml index 95ddc5e6..632ddf51 100644 --- a/crates/lean_prover/Cargo.toml +++ b/crates/lean_prover/Cargo.toml @@ -35,4 +35,5 @@ multilinear-toolkit.workspace = true poseidon_circuit.workspace = true [dev-dependencies] -xmss.workspace = true \ No newline at end of file +xmss.workspace = true +rand_chacha = "0.9.0" diff --git a/crates/lean_prover/tests/test_range_check.rs b/crates/lean_prover/tests/test_range_check.rs new file mode 100644 index 00000000..30083cff --- /dev/null +++ b/crates/lean_prover/tests/test_range_check.rs @@ -0,0 +1,159 @@ +use lean_compiler::*; +use lean_prover::{ + prove_execution::prove_execution, whir_config_builder, +}; +use lean_vm::*; +use p3_field::PrimeCharacteristicRing; +use std::collections::BTreeSet; +use rand::Rng; +use rand_chacha::ChaCha20Rng; +use rand_chacha::rand_core::SeedableRng; +use rayon::iter::{IntoParallelRefIterator, ParallelIterator}; + +const NO_VEC_RUNTIME_MEMORY: usize = 1 << 20; + +fn critical_test_cases() -> (BTreeSet<(usize, usize)>, BTreeSet<(usize, usize)>) { + let mut happy_test_cases = BTreeSet::<(usize, usize)>::new(); + let mut sad_test_cases = BTreeSet::<(usize, usize)>::new(); + + for t in 69..70 { + for v in 0..t { + if v < t { + happy_test_cases.insert((v, t)); + } else { + sad_test_cases.insert((v, t)); + } + } + //for v in 16777215..16777300 { + //if v < t { + //happy_test_cases.insert((v, t)); + //} else { + //sad_test_cases.insert((v, t)); + //} + //} + } + + (happy_test_cases, sad_test_cases) +} + +fn random_test_cases(num_test_cases: usize) -> BTreeSet<(usize, usize)> { + let v_max = 16777215 * 2; + let t_max = 65536; + let mut rng = ChaCha20Rng::seed_from_u64(0); + + let mut test_cases = BTreeSet::new(); + + for _ in 0..num_test_cases / 2 { + let t = rng.random_range(0..t_max) as usize; + let v = rng.random_range(0..v_max) as usize; + test_cases.insert((v, t)); + } + + for _ in 0..num_test_cases / 2 { + let t = rng.random_range(0..t_max) as usize; + let v = rng.random_range(0..v_max) as usize; + if v >= t { + test_cases.insert((v, t)); + } else { + test_cases.insert((t, v)); + } + } + + test_cases +} + +fn range_check_program(value: usize, max: usize) -> String { + let program = format!(r#" + const DIM = 5; + const SECOND_POINT = 2; + const SECOND_N_VARS = 7; + const COMPRESSION = 1; + const PERMUTATION = 0; + + fn main() {{ + x = 1; + y = {value}; + value = x * y; + range_check(value, {max}); + + // Need to add the following to avoid a "TODO small GKR, no packing" error + + for i in 10..50 {{ + x = malloc_vec(6); + poseidon16(i + 3, i, x, PERMUTATION); + poseidon24(i + 3, i, x + 2); + dot_product(i*2, i, (x + 3) * 8, 1); + dot_product(i*3, i + 7, (x + 4) * 8, 2); + }} + + for i in 0..1000 {{ + assert i != 1000; + }} + + return; + }} + "#); + program.to_string() +} + +fn do_test_range_check(v: usize, t: usize) { + let program_str = range_check_program(v, t); + + const SECOND_POINT: usize = 2; + const SECOND_N_VARS: usize = 7; + + let mut public_input = (0..(1 << 13) - PUBLIC_INPUT_START) + .map(F::from_usize) + .collect::>(); + + public_input[SECOND_POINT * (SECOND_N_VARS * DIMENSION).next_power_of_two() + + SECOND_N_VARS * DIMENSION + - PUBLIC_INPUT_START + ..(SECOND_POINT + 1) * (SECOND_N_VARS * DIMENSION).next_power_of_two() + - PUBLIC_INPUT_START] + .iter_mut() + .for_each(|x| *x = F::ZERO); + + let private_input = (0..1 << 13) + .map(|i| F::from_usize(i).square()) + .collect::>(); + + let bytecode = compile_program(program_str); + let _proof_data = prove_execution( + &bytecode, + (&public_input, &private_input), + whir_config_builder(), + NO_VEC_RUNTIME_MEMORY, + false, + (&vec![], &vec![]), + ); +} + +#[test] +fn test_prove_range_check_happy() { + let (happy_test_cases, _sad_test_cases) = critical_test_cases(); + println!("Running {} test cases:", happy_test_cases.len()); + //happy_test_cases.par_iter().for_each(|(v, t)| { + //do_test_range_check(*v, *t); + //}); + for (v, t) in happy_test_cases { + do_test_range_check(v, t); + } +} + +//#[test] +//fn test_prove_range_check_sad() { + //let (_happy_test_cases, sad_test_cases) = critical_test_cases(); + //for (v, t) in sad_test_cases { + //let result = std::panic::catch_unwind(|| { + //do_test_range_check(v, t); + //}); + //assert!(result.is_err(), "Expected panic for test case v={}, t={}", v, t); + //} +//} + +#[test] +#[should_panic] +fn test_prove_range_check_sad_1() { + do_test_range_check(0, 0); +} diff --git a/crates/lean_prover/witness_generation/src/execution_trace.rs b/crates/lean_prover/witness_generation/src/execution_trace.rs index 84773333..f5cb4f79 100644 --- a/crates/lean_prover/witness_generation/src/execution_trace.rs +++ b/crates/lean_prover/witness_generation/src/execution_trace.rs @@ -55,6 +55,7 @@ pub fn get_execution_trace( .zip(execution_result.fps.par_iter()) .for_each(|((trace_row, &pc), &fp)| { let instruction = &bytecode.instructions[pc]; + //println!("instruction: {}", instruction); let field_repr = field_representation(instruction); let mut addr_a = F::ZERO; diff --git a/crates/lean_prover/witness_generation/src/instruction_encoder.rs b/crates/lean_prover/witness_generation/src/instruction_encoder.rs index e8700705..ec551235 100644 --- a/crates/lean_prover/witness_generation/src/instruction_encoder.rs +++ b/crates/lean_prover/witness_generation/src/instruction_encoder.rs @@ -29,6 +29,7 @@ pub fn field_representation(instr: &Instruction) -> [F; N_INSTRUCTION_COLUMNS] { shift_0, shift_1, res, + for_range_check: _, } => { fields[3] = F::ZERO; // flag_A = 0 fields[0] = F::from_usize(*shift_0); diff --git a/crates/lean_vm/src/isa/hint.rs b/crates/lean_vm/src/isa/hint.rs index 7b4f1c70..814438c6 100644 --- a/crates/lean_vm/src/isa/hint.rs +++ b/crates/lean_vm/src/isa/hint.rs @@ -1,7 +1,7 @@ use crate::core::{F, LOG_VECTOR_LEN, Label, SourceLineNumber, VECTOR_LEN}; use crate::diagnostics::RunnerError; use crate::execution::{ExecutionHistory, Memory}; -use crate::isa::operands::MemOrConstant; +use crate::isa::operands::{MemOrConstant, MemOrFp}; use p3_field::{Field, PrimeCharacteristicRing}; use std::fmt::{Display, Formatter}; use utils::{ToUsize, pretty_integer}; @@ -65,6 +65,12 @@ pub enum Hint { }, /// Jump destination label (for debugging purposes) Label { label: Label }, + + /// Range check + RangeCheck { + value: MemOrFp, + max: MemOrConstant, + }, } /// Execution state for hint processing @@ -193,6 +199,7 @@ impl Hint { *ctx.cpu_cycles_before_new_line = 0; } Self::Label { .. } => {} + Self::RangeCheck { .. } => {} } Ok(()) } @@ -266,6 +273,9 @@ impl Display for Hint { Self::Label { label } => { write!(f, "label: {label}") } + Self::RangeCheck { value, max } => { + write!(f, "range_check({value}, {max})") + } } } } diff --git a/crates/lean_vm/src/isa/instruction.rs b/crates/lean_vm/src/isa/instruction.rs index 7810ff2d..138acd6c 100644 --- a/crates/lean_vm/src/isa/instruction.rs +++ b/crates/lean_vm/src/isa/instruction.rs @@ -39,6 +39,7 @@ pub enum Instruction { shift_1: usize, /// Result destination (can be memory, frame pointer, or constant) res: MemOrFpOrConstant, + for_range_check: bool, }, /// Conditional jump instruction for control flow @@ -176,12 +177,23 @@ impl Instruction { shift_0, shift_1, res, + for_range_check, } => { if res.is_value_unknown(ctx.memory, *ctx.fp) { let memory_address_res = res.memory_address(*ctx.fp)?; let ptr = ctx.memory.get(*ctx.fp + shift_0)?; - let value = ctx.memory.get(ptr.to_usize() + shift_1)?; - ctx.memory.set(memory_address_res, value)?; + + if *for_range_check { + // if for_range_check, ignore the UndefinedMemory error from get() + let value = ctx.memory.get(ptr.to_usize() + shift_1); + if value.is_ok() { + ctx.memory.set(memory_address_res, value.unwrap())?; + } + } else { + // else, bubble it up + let value = ctx.memory.get(ptr.to_usize() + shift_1)?; + ctx.memory.set(memory_address_res, value)?; + } } else { let value = res.read_value(ctx.memory, *ctx.fp)?; let ptr = ctx.memory.get(*ctx.fp + shift_0)?; @@ -424,8 +436,9 @@ impl Display for Instruction { shift_0, shift_1, res, + for_range_check, } => { - write!(f, "{res} = m[m[fp + {shift_0}] + {shift_1}]") + write!(f, "{res} = m[m[fp + {shift_0}] + {shift_1}] for_range_check: {for_range_check}") } Self::DotProductExtensionExtension { arg0, From 08fea21a44cffea5699da8f39be0f66fb8e72b76 Mon Sep 17 00:00:00 2001 From: Koh Wei Jie Date: Tue, 28 Oct 2025 20:36:15 -0700 Subject: [PATCH 02/10] added clarifying comments and more tests --- crates/lean_prover/tests/test_range_check.rs | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/crates/lean_prover/tests/test_range_check.rs b/crates/lean_prover/tests/test_range_check.rs index 30083cff..6855d428 100644 --- a/crates/lean_prover/tests/test_range_check.rs +++ b/crates/lean_prover/tests/test_range_check.rs @@ -157,3 +157,21 @@ fn test_prove_range_check_happy() { fn test_prove_range_check_sad_1() { do_test_range_check(0, 0); } + +#[test] +#[should_panic] +fn test_prove_range_check_sad_2() { + do_test_range_check(1, 0); +} + +#[test] +#[should_panic] +fn test_prove_range_check_sad_3() { + do_test_range_check(2, 1); +} + +#[test] +#[should_panic] +fn test_prove_range_check_sad_4() { + do_test_range_check(69, 65); +} From c448cd3eae29e2f62bda03e9416c68561079dff1 Mon Sep 17 00:00:00 2001 From: Koh Wei Jie Date: Fri, 31 Oct 2025 11:28:19 -0700 Subject: [PATCH 03/10] completed range check impl --- Cargo.lock | 1 + .../src/b_compile_intermediate.rs | 41 ++--- crates/lean_prover/Cargo.toml | 1 + crates/lean_prover/src/prove_execution.rs | 21 ++- crates/lean_prover/tests/hash_chain.rs | 1 + crates/lean_prover/tests/test_range_check.rs | 169 +++++++----------- crates/lean_prover/tests/test_zkvm.rs | 1 + .../witness_generation/src/execution_trace.rs | 6 +- crates/utils/src/misc.rs | 2 +- 9 files changed, 108 insertions(+), 135 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bed5b709..cffa761e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -405,6 +405,7 @@ dependencies = [ "rand", "rand_chacha", "rayon", + "thiserror", "tracing", "utils", "vm_air", diff --git a/crates/lean_compiler/src/b_compile_intermediate.rs b/crates/lean_compiler/src/b_compile_intermediate.rs index 899eadf4..8dd7a741 100644 --- a/crates/lean_compiler/src/b_compile_intermediate.rs +++ b/crates/lean_compiler/src/b_compile_intermediate.rs @@ -615,6 +615,8 @@ fn compile_lines( }); } SimpleLine::RangeCheck { value, max } => { + // x is the fp offset of the memory cell which contains the value + // i.e. m[fp + x] contains value let x = match IntermediateValue::from_simple_expr(value, compiler) { IntermediateValue::MemoryAfterFp { offset } => offset.naive_eval().unwrap(), value::IntermediateValue::Fp => F::ZERO, @@ -636,17 +638,14 @@ fn compile_lines( for_range_check: true, }; - - // Step 2: ADD: m[m[fp + x]] + m[fp + j] == (t-1) - // m[fp + j] == t - 1 - value - // + // Step 2: ADD: m[fp + x] + m[fp + j] == (t-1) // m[fp + j] == t - 1 - m[fp + x] - let q = t - F::ONE; + // Uses constraint solving to store t - 1 - m[fp + x] in m[fp + j] let step_2 = IntermediateInstruction::Computation { operation: Operation::Add, arg_a: IntermediateValue::MemoryAfterFp { offset: x.to_usize().into() }, - arg_c: IntermediateValue::MemoryAfterFp { offset: aux_j.into() }, // solve - res: IntermediateValue::Constant(q.to_usize().into()), // t - 1 + arg_c: IntermediateValue::MemoryAfterFp { offset: aux_j.into() }, + res: IntermediateValue::Constant((t - F::ONE).to_usize().into()), }; // Step 3: DEREF: m[fp + k] == m[m[fp + j]] @@ -657,20 +656,22 @@ fn compile_lines( for_range_check: true, }; + // Insert the instructions + instructions.extend_from_slice( + &[ + // This is just the RangeCheck hint which does nothing + IntermediateInstruction::RangeCheck { + value: IntermediateValue::from_simple_expr(value, compiler), + max: max.clone(), + }, + // These are the steps that effectuate the range check + step_1, + step_2, + step_3, + ] + ); - // TODO: handle undefined memory access error - - //println!("aux_i: {}; {:?}", aux_i, step_1); - //println!("aux_j: {}; {:?}", aux_j, step_2); - //println!("aux_k: {}; {:?}", aux_k, step_3); - - instructions.push(IntermediateInstruction::RangeCheck { - value: IntermediateValue::from_simple_expr(value, compiler), - max: max.clone(), - }); - instructions.push(step_1); - instructions.push(step_2); - instructions.push(step_3); + // Increase the stack size by 3 as we used 3 aux variables compiler.stack_size += 3; } } diff --git a/crates/lean_prover/Cargo.toml b/crates/lean_prover/Cargo.toml index 632ddf51..bcf4e7ec 100644 --- a/crates/lean_prover/Cargo.toml +++ b/crates/lean_prover/Cargo.toml @@ -33,6 +33,7 @@ witness_generation.workspace = true vm_air.workspace = true multilinear-toolkit.workspace = true poseidon_circuit.workspace = true +thiserror.workspace = true [dev-dependencies] xmss.workspace = true diff --git a/crates/lean_prover/src/prove_execution.rs b/crates/lean_prover/src/prove_execution.rs index 45460b32..4acf2349 100644 --- a/crates/lean_prover/src/prove_execution.rs +++ b/crates/lean_prover/src/prove_execution.rs @@ -22,6 +22,14 @@ use whir_p3::{ }; use xmss::{Poseidon16History, Poseidon24History}; +use thiserror::Error; + +#[derive(Error, Debug)] +pub enum ProveExecutionError { + #[error("MemoryError")] + MemoryError, +} + pub fn prove_execution( bytecode: &Bytecode, (public_input, private_input): (&[F], &[F]), @@ -29,7 +37,7 @@ pub fn prove_execution( no_vec_runtime_memory: usize, // size of the "non-vectorized" runtime memory vm_profiler: bool, (poseidons_16_precomputed, poseidons_24_precomputed): (&Poseidon16History, &Poseidon24History), -) -> (Vec>, usize, String) { +) -> Result<(Vec>, usize, String), ProveExecutionError> { let mut exec_summary = String::new(); let ExecutionTrace { full_trace, @@ -843,9 +851,14 @@ pub fn prove_execution( &mut base_memory_poly_eq_point, memory_poly_eq_point_alpha.square(), ); + + let memory_len = memory.len(); + if base_memory_indexes.iter().max().unwrap().to_usize() >= memory_len { + return Err(ProveExecutionError::MemoryError); + } let base_memory_pushforward = compute_pushforward( &base_memory_indexes, - memory.len(), + memory_len, &base_memory_poly_eq_point, ); @@ -1188,9 +1201,9 @@ pub fn prove_execution( &packed_pcs_witness_extension.packed_polynomial.by_ref(), ); - ( + Ok(( prover_state.proof_data().to_vec(), prover_state.proof_size(), exec_summary, - ) + )) } diff --git a/crates/lean_prover/tests/hash_chain.rs b/crates/lean_prover/tests/hash_chain.rs index d4a9f874..505817b4 100644 --- a/crates/lean_prover/tests/hash_chain.rs +++ b/crates/lean_prover/tests/hash_chain.rs @@ -81,6 +81,7 @@ fn benchmark_poseidon_chain() { false, (&vec![], &vec![]), // TODO poseidons precomputed ) + .unwrap() .0; let vm_time = time.elapsed(); verify_execution(&bytecode, &public_input, proof_data, whir_config_builder()).unwrap(); diff --git a/crates/lean_prover/tests/test_range_check.rs b/crates/lean_prover/tests/test_range_check.rs index 6855d428..041e4e82 100644 --- a/crates/lean_prover/tests/test_range_check.rs +++ b/crates/lean_prover/tests/test_range_check.rs @@ -1,39 +1,43 @@ -use lean_compiler::*; -use lean_prover::{ - prove_execution::prove_execution, whir_config_builder, -}; -use lean_vm::*; +use lean_compiler::compile_program; +use lean_vm::{F, DIMENSION, PUBLIC_INPUT_START}; +use lean_prover::{whir_config_builder, prove_execution::prove_execution}; +use whir_p3::WhirConfigBuilder; use p3_field::PrimeCharacteristicRing; use std::collections::BTreeSet; use rand::Rng; use rand_chacha::ChaCha20Rng; use rand_chacha::rand_core::SeedableRng; -use rayon::iter::{IntoParallelRefIterator, ParallelIterator}; const NO_VEC_RUNTIME_MEMORY: usize = 1 << 20; -fn critical_test_cases() -> (BTreeSet<(usize, usize)>, BTreeSet<(usize, usize)>) { - let mut happy_test_cases = BTreeSet::<(usize, usize)>::new(); - let mut sad_test_cases = BTreeSet::<(usize, usize)>::new(); - - for t in 69..70 { - for v in 0..t { - if v < t { - happy_test_cases.insert((v, t)); - } else { - sad_test_cases.insert((v, t)); - } - } - //for v in 16777215..16777300 { - //if v < t { - //happy_test_cases.insert((v, t)); - //} else { - //sad_test_cases.insert((v, t)); - //} - //} - } +fn range_check_program(value: usize, max: usize) -> String { + let program = format!(r#" + fn func() {{ + x = 1; + y = {value}; + value = x * y; + range_check(value, {max}); + return; + }} + + fn main() {{ + x = 1; + y = {value}; + value = x * y; + range_check(value, {max}); - (happy_test_cases, sad_test_cases) + func(); + + if 0 == 0 {{ + a = 1; + b = {value}; + c = a * b; + range_check(c, {max}); + }} + return; + }} + "#); + program.to_string() } fn random_test_cases(num_test_cases: usize) -> BTreeSet<(usize, usize)> { @@ -62,43 +66,7 @@ fn random_test_cases(num_test_cases: usize) -> BTreeSet<(usize, usize)> { test_cases } -fn range_check_program(value: usize, max: usize) -> String { - let program = format!(r#" - const DIM = 5; - const SECOND_POINT = 2; - const SECOND_N_VARS = 7; - const COMPRESSION = 1; - const PERMUTATION = 0; - - fn main() {{ - x = 1; - y = {value}; - value = x * y; - range_check(value, {max}); - - // Need to add the following to avoid a "TODO small GKR, no packing" error - - for i in 10..50 {{ - x = malloc_vec(6); - poseidon16(i + 3, i, x, PERMUTATION); - poseidon24(i + 3, i, x + 2); - dot_product(i*2, i, (x + 3) * 8, 1); - dot_product(i*3, i + 7, (x + 4) * 8, 2); - }} - - for i in 0..1000 {{ - assert i != 1000; - }} - - return; - }} - "#); - program.to_string() -} - -fn do_test_range_check(v: usize, t: usize) { - let program_str = range_check_program(v, t); - +fn prepare_inputs() -> (Vec, Vec) { const SECOND_POINT: usize = 2; const SECOND_N_VARS: usize = 7; @@ -117,61 +85,48 @@ fn do_test_range_check(v: usize, t: usize) { let private_input = (0..1 << 13) .map(|i| F::from_usize(i).square()) .collect::>(); + + (public_input, private_input) +} + +fn do_test_range_check( + v: usize, + t: usize, + whir_config_builder: &WhirConfigBuilder, + public_input: &Vec, + private_input: &Vec +) { + let program_str = range_check_program(v, t); let bytecode = compile_program(program_str); - let _proof_data = prove_execution( + + let r = prove_execution( &bytecode, - (&public_input, &private_input), - whir_config_builder(), + (public_input, private_input), + whir_config_builder.clone(), NO_VEC_RUNTIME_MEMORY, false, (&vec![], &vec![]), ); -} - -#[test] -fn test_prove_range_check_happy() { - let (happy_test_cases, _sad_test_cases) = critical_test_cases(); - println!("Running {} test cases:", happy_test_cases.len()); - //happy_test_cases.par_iter().for_each(|(v, t)| { - //do_test_range_check(*v, *t); - //}); - for (v, t) in happy_test_cases { - do_test_range_check(v, t); + + if v < t { + assert!(r.is_ok(), "Proof generation should work for v < t"); + } else { + assert!(r.is_err(), "Proof generation should fail for v >= t"); } -} - -//#[test] -//fn test_prove_range_check_sad() { - //let (_happy_test_cases, sad_test_cases) = critical_test_cases(); - //for (v, t) in sad_test_cases { - //let result = std::panic::catch_unwind(|| { - //do_test_range_check(v, t); - //}); - //assert!(result.is_err(), "Expected panic for test case v={}, t={}", v, t); - //} -//} -#[test] -#[should_panic] -fn test_prove_range_check_sad_1() { - do_test_range_check(0, 0); } #[test] -#[should_panic] -fn test_prove_range_check_sad_2() { - do_test_range_check(1, 0); -} +fn test_range_check() { + let (public_input, private_input) = prepare_inputs(); + let whir_config_builder = whir_config_builder(); -#[test] -#[should_panic] -fn test_prove_range_check_sad_3() { - do_test_range_check(2, 1); -} + let test_cases = random_test_cases(500); -#[test] -#[should_panic] -fn test_prove_range_check_sad_4() { - do_test_range_check(69, 65); + println!("Running {} random test cases", test_cases.len()); + + for (v, t) in test_cases { + do_test_range_check(v, t, &whir_config_builder, &public_input, &private_input); + } } diff --git a/crates/lean_prover/tests/test_zkvm.rs b/crates/lean_prover/tests/test_zkvm.rs index bdcd08eb..acf4ba31 100644 --- a/crates/lean_prover/tests/test_zkvm.rs +++ b/crates/lean_prover/tests/test_zkvm.rs @@ -109,6 +109,7 @@ fn test_zk_vm_helper(program_str: &str) { false, (&vec![], &vec![]), ) + .unwrap() .0; verify_execution(&bytecode, &public_input, proof_data, whir_config_builder()).unwrap(); } diff --git a/crates/lean_prover/witness_generation/src/execution_trace.rs b/crates/lean_prover/witness_generation/src/execution_trace.rs index f5cb4f79..4e3e5cb6 100644 --- a/crates/lean_prover/witness_generation/src/execution_trace.rs +++ b/crates/lean_prover/witness_generation/src/execution_trace.rs @@ -63,13 +63,13 @@ pub fn get_execution_trace( // flag_a == 0 addr_a = F::from_usize(fp) + field_repr[0]; // fp + operand_a } - let value_a = memory.0[addr_a.to_usize()].unwrap(); + let value_a = memory.get(addr_a.to_usize()).unwrap_or(F::ZERO); let mut addr_b = F::ZERO; if field_repr[4].is_zero() { // flag_b == 0 addr_b = F::from_usize(fp) + field_repr[1]; // fp + operand_b } - let value_b = memory.0[addr_b.to_usize()].unwrap(); + let value_b = memory.get(addr_b.to_usize()).unwrap_or(F::ZERO); let mut addr_c = F::ZERO; if field_repr[5].is_zero() { @@ -80,7 +80,7 @@ pub fn get_execution_trace( assert_eq!(field_repr[2], operand_c); // debug purpose addr_c = value_a + operand_c; } - let value_c = memory.0[addr_c.to_usize()].unwrap(); + let value_c = memory.get(addr_c.to_usize()).unwrap_or(F::ZERO); for (j, field) in field_repr.iter().enumerate() { *trace_row[j] = *field; diff --git a/crates/utils/src/misc.rs b/crates/utils/src/misc.rs index 02d8d257..ec82aa6a 100644 --- a/crates/utils/src/misc.rs +++ b/crates/utils/src/misc.rs @@ -111,7 +111,7 @@ pub fn transpose( width: usize, column_extra_capacity: usize, ) -> Vec> { - assert!((matrix.len().is_multiple_of(width))); + assert!(matrix.len().is_multiple_of(width)); let height = matrix.len() / width; let res = vec![ { From 0c8b13f59d4fc4f9436898b3d7d455d9f27d8a28 Mon Sep 17 00:00:00 2001 From: Koh Wei Jie Date: Sun, 2 Nov 2025 11:32:26 -0800 Subject: [PATCH 04/10] slight optimisation to check base_memory_indexes in the prover --- crates/lean_prover/src/prove_execution.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/crates/lean_prover/src/prove_execution.rs b/crates/lean_prover/src/prove_execution.rs index 4acf2349..8f0f188f 100644 --- a/crates/lean_prover/src/prove_execution.rs +++ b/crates/lean_prover/src/prove_execution.rs @@ -853,9 +853,10 @@ pub fn prove_execution( ); let memory_len = memory.len(); - if base_memory_indexes.iter().max().unwrap().to_usize() >= memory_len { + if base_memory_indexes.iter().any(|i| i.to_usize() >= memory_len) { return Err(ProveExecutionError::MemoryError); } + let base_memory_pushforward = compute_pushforward( &base_memory_indexes, memory_len, From 94c46b3d45eb3a77c1ce9a5346073d637bddf6c8 Mon Sep 17 00:00:00 2001 From: Koh Wei Jie Date: Sun, 2 Nov 2025 11:40:08 -0800 Subject: [PATCH 05/10] cargo fmt --- .../src/b_compile_intermediate.rs | 44 +++++++++++-------- crates/lean_compiler/src/c_compile_final.rs | 2 +- crates/lean_compiler/src/ir/instruction.rs | 5 ++- crates/lean_prover/src/prove_execution.rs | 12 ++--- crates/lean_prover/tests/test_range_check.rs | 25 ++++++----- crates/lean_vm/src/isa/hint.rs | 5 +-- crates/lean_vm/src/isa/instruction.rs | 7 ++- 7 files changed, 55 insertions(+), 45 deletions(-) diff --git a/crates/lean_compiler/src/b_compile_intermediate.rs b/crates/lean_compiler/src/b_compile_intermediate.rs index 8dd7a741..a9e544c0 100644 --- a/crates/lean_compiler/src/b_compile_intermediate.rs +++ b/crates/lean_compiler/src/b_compile_intermediate.rs @@ -634,17 +634,23 @@ fn compile_lines( let step_1 = IntermediateInstruction::Deref { shift_0: ConstExpression::scalar(x.to_usize()), shift_1: ConstExpression::from(0), - res: IntermediaryMemOrFpOrConstant::MemoryAfterFp { offset: aux_i.into() }, + res: IntermediaryMemOrFpOrConstant::MemoryAfterFp { + offset: aux_i.into(), + }, for_range_check: true, }; - // Step 2: ADD: m[fp + x] + m[fp + j] == (t-1) + // Step 2: ADD: m[fp + x] + m[fp + j] == (t-1) // m[fp + j] == t - 1 - m[fp + x] // Uses constraint solving to store t - 1 - m[fp + x] in m[fp + j] let step_2 = IntermediateInstruction::Computation { operation: Operation::Add, - arg_a: IntermediateValue::MemoryAfterFp { offset: x.to_usize().into() }, - arg_c: IntermediateValue::MemoryAfterFp { offset: aux_j.into() }, + arg_a: IntermediateValue::MemoryAfterFp { + offset: x.to_usize().into(), + }, + arg_c: IntermediateValue::MemoryAfterFp { + offset: aux_j.into(), + }, res: IntermediateValue::Constant((t - F::ONE).to_usize().into()), }; @@ -652,24 +658,24 @@ fn compile_lines( let step_3 = IntermediateInstruction::Deref { shift_0: ConstExpression::scalar(aux_j), shift_1: ConstExpression::from(0), - res: IntermediaryMemOrFpOrConstant::MemoryAfterFp { offset: aux_k.into() }, + res: IntermediaryMemOrFpOrConstant::MemoryAfterFp { + offset: aux_k.into(), + }, for_range_check: true, }; // Insert the instructions - instructions.extend_from_slice( - &[ - // This is just the RangeCheck hint which does nothing - IntermediateInstruction::RangeCheck { - value: IntermediateValue::from_simple_expr(value, compiler), - max: max.clone(), - }, - // These are the steps that effectuate the range check - step_1, - step_2, - step_3, - ] - ); + instructions.extend_from_slice(&[ + // This is just the RangeCheck hint which does nothing + IntermediateInstruction::RangeCheck { + value: IntermediateValue::from_simple_expr(value, compiler), + max: max.clone(), + }, + // These are the steps that effectuate the range check + step_1, + step_2, + step_3, + ]); // Increase the stack size by 3 as we used 3 aux variables compiler.stack_size += 3; @@ -876,7 +882,7 @@ fn find_internal_vars(lines: &[SimpleLine]) -> BTreeSet { | SimpleLine::Print { .. } | SimpleLine::FunctionRet { .. } | SimpleLine::Precompile { .. } - | SimpleLine::LocationReport { .. } + | SimpleLine::LocationReport { .. } | SimpleLine::RangeCheck { .. } => {} } } diff --git a/crates/lean_compiler/src/c_compile_final.rs b/crates/lean_compiler/src/c_compile_final.rs index 30d11f1f..e7c60e61 100644 --- a/crates/lean_compiler/src/c_compile_final.rs +++ b/crates/lean_compiler/src/c_compile_final.rs @@ -23,7 +23,7 @@ impl IntermediateInstruction { | Self::Poseidon2_16 { .. } | Self::Poseidon2_24 { .. } | Self::DotProduct { .. } - | Self::MultilinearEval { .. } => false + | Self::MultilinearEval { .. } => false, } } } diff --git a/crates/lean_compiler/src/ir/instruction.rs b/crates/lean_compiler/src/ir/instruction.rs index d2ff001a..a1922112 100644 --- a/crates/lean_compiler/src/ir/instruction.rs +++ b/crates/lean_compiler/src/ir/instruction.rs @@ -158,7 +158,10 @@ impl Display for IntermediateInstruction { shift_1, res, for_range_check, - } => write!(f, "{res} = m[m[fp + {shift_0}] + {shift_1}] for_range_check: {for_range_check}"), + } => write!( + f, + "{res} = m[m[fp + {shift_0}] + {shift_1}] for_range_check: {for_range_check}" + ), Self::Panic => write!(f, "panic"), Self::Jump { dest, updated_fp } => { if let Some(fp) = updated_fp { diff --git a/crates/lean_prover/src/prove_execution.rs b/crates/lean_prover/src/prove_execution.rs index 8f0f188f..cf31eedf 100644 --- a/crates/lean_prover/src/prove_execution.rs +++ b/crates/lean_prover/src/prove_execution.rs @@ -853,15 +853,15 @@ pub fn prove_execution( ); let memory_len = memory.len(); - if base_memory_indexes.iter().any(|i| i.to_usize() >= memory_len) { + if base_memory_indexes + .iter() + .any(|i| i.to_usize() >= memory_len) + { return Err(ProveExecutionError::MemoryError); } - let base_memory_pushforward = compute_pushforward( - &base_memory_indexes, - memory_len, - &base_memory_poly_eq_point, - ); + let base_memory_pushforward = + compute_pushforward(&base_memory_indexes, memory_len, &base_memory_poly_eq_point); // 2nd Commitment let extension_pols = vec![ diff --git a/crates/lean_prover/tests/test_range_check.rs b/crates/lean_prover/tests/test_range_check.rs index 041e4e82..41da73e7 100644 --- a/crates/lean_prover/tests/test_range_check.rs +++ b/crates/lean_prover/tests/test_range_check.rs @@ -1,17 +1,18 @@ use lean_compiler::compile_program; -use lean_vm::{F, DIMENSION, PUBLIC_INPUT_START}; -use lean_prover::{whir_config_builder, prove_execution::prove_execution}; -use whir_p3::WhirConfigBuilder; +use lean_prover::{prove_execution::prove_execution, whir_config_builder}; +use lean_vm::{DIMENSION, F, PUBLIC_INPUT_START}; use p3_field::PrimeCharacteristicRing; -use std::collections::BTreeSet; use rand::Rng; use rand_chacha::ChaCha20Rng; use rand_chacha::rand_core::SeedableRng; +use std::collections::BTreeSet; +use whir_p3::WhirConfigBuilder; const NO_VEC_RUNTIME_MEMORY: usize = 1 << 20; fn range_check_program(value: usize, max: usize) -> String { - let program = format!(r#" + let program = format!( + r#" fn func() {{ x = 1; y = {value}; @@ -36,8 +37,9 @@ fn range_check_program(value: usize, max: usize) -> String { }} return; }} - "#); - program.to_string() + "# + ); + program.to_string() } fn random_test_cases(num_test_cases: usize) -> BTreeSet<(usize, usize)> { @@ -85,16 +87,16 @@ fn prepare_inputs() -> (Vec, Vec) { let private_input = (0..1 << 13) .map(|i| F::from_usize(i).square()) .collect::>(); - + (public_input, private_input) } fn do_test_range_check( v: usize, - t: usize, + t: usize, whir_config_builder: &WhirConfigBuilder, public_input: &Vec, - private_input: &Vec + private_input: &Vec, ) { let program_str = range_check_program(v, t); @@ -108,13 +110,12 @@ fn do_test_range_check( false, (&vec![], &vec![]), ); - + if v < t { assert!(r.is_ok(), "Proof generation should work for v < t"); } else { assert!(r.is_err(), "Proof generation should fail for v >= t"); } - } #[test] diff --git a/crates/lean_vm/src/isa/hint.rs b/crates/lean_vm/src/isa/hint.rs index 814438c6..ea5e5493 100644 --- a/crates/lean_vm/src/isa/hint.rs +++ b/crates/lean_vm/src/isa/hint.rs @@ -67,10 +67,7 @@ pub enum Hint { Label { label: Label }, /// Range check - RangeCheck { - value: MemOrFp, - max: MemOrConstant, - }, + RangeCheck { value: MemOrFp, max: MemOrConstant }, } /// Execution state for hint processing diff --git a/crates/lean_vm/src/isa/instruction.rs b/crates/lean_vm/src/isa/instruction.rs index 138acd6c..81cb91ac 100644 --- a/crates/lean_vm/src/isa/instruction.rs +++ b/crates/lean_vm/src/isa/instruction.rs @@ -184,7 +184,7 @@ impl Instruction { let ptr = ctx.memory.get(*ctx.fp + shift_0)?; if *for_range_check { - // if for_range_check, ignore the UndefinedMemory error from get() + // if for_range_check, ignore the UndefinedMemory error from get() let value = ctx.memory.get(ptr.to_usize() + shift_1); if value.is_ok() { ctx.memory.set(memory_address_res, value.unwrap())?; @@ -438,7 +438,10 @@ impl Display for Instruction { res, for_range_check, } => { - write!(f, "{res} = m[m[fp + {shift_0}] + {shift_1}] for_range_check: {for_range_check}") + write!( + f, + "{res} = m[m[fp + {shift_0}] + {shift_1}] for_range_check: {for_range_check}" + ) } Self::DotProductExtensionExtension { arg0, From ccef7874826e6bae05dab7eca2c2719e9d60e5cd Mon Sep 17 00:00:00 2001 From: Tom Wambsgans Date: Tue, 4 Nov 2025 12:22:26 +0400 Subject: [PATCH 06/10] clippy --- crates/lean_compiler/src/b_compile_intermediate.rs | 6 +++--- crates/lean_prover/tests/test_range_check.rs | 12 ++++++------ 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/crates/lean_compiler/src/b_compile_intermediate.rs b/crates/lean_compiler/src/b_compile_intermediate.rs index 01ae084a..2a03655d 100644 --- a/crates/lean_compiler/src/b_compile_intermediate.rs +++ b/crates/lean_compiler/src/b_compile_intermediate.rs @@ -648,9 +648,9 @@ fn compile_lines( }; let t = max.naive_eval().unwrap(); - let aux_i: usize = compiler.stack_size + 0; - let aux_j: usize = compiler.stack_size + 1; - let aux_k: usize = compiler.stack_size + 2; + let aux_i = compiler.stack_size; + let aux_j = compiler.stack_size + 1; + let aux_k = compiler.stack_size + 2; // Step 1: DEREF: m[fp + i] == m[m[fp + x]] // DEREF: m[fp + i] == m[value] diff --git a/crates/lean_prover/tests/test_range_check.rs b/crates/lean_prover/tests/test_range_check.rs index eb4fca6f..864a91fd 100644 --- a/crates/lean_prover/tests/test_range_check.rs +++ b/crates/lean_prover/tests/test_range_check.rs @@ -49,11 +49,11 @@ fn random_test_cases(num_test_cases: usize, valid: bool) -> BTreeSet<(usize, usi let mut test_cases = BTreeSet::new(); for _ in 0..num_test_cases { - let t = rng.random_range(1..t_max) as usize; + let t = rng.random_range(1..t_max); let v = if valid { - rng.random_range(0..t) as usize + rng.random_range(0..t) } else { - rng.random_range(t..1 << 31) as usize + rng.random_range(t..1 << 31) }; test_cases.insert((v, t)); } @@ -88,8 +88,8 @@ fn do_test_range_check( v: usize, t: usize, whir_config_builder: &WhirConfigBuilder, - public_input: &Vec, - private_input: &Vec, + public_input: &[F], + private_input: &[F], ) { let program_str = range_check_program(v, t); @@ -105,7 +105,7 @@ fn do_test_range_check( ); verify_execution( &bytecode, - &public_input, + public_input, proof_data, whir_config_builder.clone(), ) From a2b25a2cdbf4471d8e91e5b88c2ddb36ae3316aa Mon Sep 17 00:00:00 2001 From: Tom Wambsgans Date: Tue, 4 Nov 2025 12:35:46 +0400 Subject: [PATCH 07/10] clippy --- crates/lean_vm/src/isa/instruction.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/lean_vm/src/isa/instruction.rs b/crates/lean_vm/src/isa/instruction.rs index 81cb91ac..3b326e1d 100644 --- a/crates/lean_vm/src/isa/instruction.rs +++ b/crates/lean_vm/src/isa/instruction.rs @@ -186,8 +186,8 @@ impl Instruction { if *for_range_check { // if for_range_check, ignore the UndefinedMemory error from get() let value = ctx.memory.get(ptr.to_usize() + shift_1); - if value.is_ok() { - ctx.memory.set(memory_address_res, value.unwrap())?; + if let Ok(value) = value { + ctx.memory.set(memory_address_res, value)?; } } else { // else, bubble it up From 404a38bd7867832d63b2a52bfa044cd3a33e4f39 Mon Sep 17 00:00:00 2001 From: Tom Wambsgans Date: Wed, 12 Nov 2025 13:00:29 +0400 Subject: [PATCH 08/10] add a failing test --- crates/lean_prover/tests/test_range_check.rs | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/crates/lean_prover/tests/test_range_check.rs b/crates/lean_prover/tests/test_range_check.rs index 864a91fd..9804a9fb 100644 --- a/crates/lean_prover/tests/test_range_check.rs +++ b/crates/lean_prover/tests/test_range_check.rs @@ -135,3 +135,17 @@ fn test_range_check(num_test_cases: usize, valid: bool) { do_test_range_check(v, t, &whir_config_builder, &public_input, &private_input); } } + +// BUG HERE +#[test] +fn failing_test() { + let (public_input, private_input) = prepare_inputs(); + let whir_config_builder = whir_config_builder(); + do_test_range_check( + 16390, + 1 << 16, + &whir_config_builder, + &public_input, + &private_input, + ); +} From 119fcc0fb413e0f2309aa0a93a9d21775c90f491 Mon Sep 17 00:00:00 2001 From: Koh Wei Jie Date: Sun, 16 Nov 2025 10:00:51 -0800 Subject: [PATCH 09/10] Fixed range checks using prover hints --- crates/lean_prover/src/prove_execution.rs | 9 +++++ crates/lean_prover/tests/test_range_check.rs | 37 +++++++++++++++++--- crates/lean_vm/src/diagnostics/error.rs | 2 ++ crates/lean_vm/src/execution/runner.rs | 3 ++ crates/lean_vm/src/isa/instruction.rs | 13 +++++-- 5 files changed, 57 insertions(+), 7 deletions(-) diff --git a/crates/lean_prover/src/prove_execution.rs b/crates/lean_prover/src/prove_execution.rs index 528db386..60dce7e1 100644 --- a/crates/lean_prover/src/prove_execution.rs +++ b/crates/lean_prover/src/prove_execution.rs @@ -54,6 +54,15 @@ pub fn prove_execution( (poseidons_16_precomputed, poseidons_24_precomputed), ) }); + + // Fill undefined memory cells used for range checks + for (dest_ptr, src_ptr) in &execution_result.range_check_cells_to_fill { + if execution_result.memory.get(*dest_ptr).is_err() { + let value = execution_result.memory.get(*src_ptr).unwrap_or(F::ZERO); + execution_result.memory.set(*dest_ptr, value).unwrap(); + } + } + exec_summary = std::mem::take(&mut execution_result.summary); info_span!("Building execution trace") .in_scope(|| get_execution_trace(bytecode, execution_result)) diff --git a/crates/lean_prover/tests/test_range_check.rs b/crates/lean_prover/tests/test_range_check.rs index 864a91fd..06a41dae 100644 --- a/crates/lean_prover/tests/test_range_check.rs +++ b/crates/lean_prover/tests/test_range_check.rs @@ -48,13 +48,14 @@ fn random_test_cases(num_test_cases: usize, valid: bool) -> BTreeSet<(usize, usi let mut test_cases = BTreeSet::new(); - for _ in 0..num_test_cases { + while test_cases.len() < num_test_cases { let t = rng.random_range(1..t_max); let v = if valid { rng.random_range(0..t) } else { rng.random_range(t..1 << 31) }; + test_cases.insert((v, t)); } @@ -114,16 +115,16 @@ fn do_test_range_check( #[test] fn test_range_check_valid() { - test_range_check(10, true); + test_range_check_random(100, true); } #[test] #[should_panic] fn test_range_check_invalid() { - test_range_check(1, false); + test_range_check_random(1, false); } -fn test_range_check(num_test_cases: usize, valid: bool) { +fn test_range_check_random(num_test_cases: usize, valid: bool) { let (public_input, private_input) = prepare_inputs(); let whir_config_builder = whir_config_builder(); @@ -132,6 +133,34 @@ fn test_range_check(num_test_cases: usize, valid: bool) { println!("Running {} random test cases", test_cases.len()); for (v, t) in test_cases { + println!("v: {v}; t: {t}"); do_test_range_check(v, t, &whir_config_builder, &public_input, &private_input); } } + +#[test] +fn test_range_check_valid_1() { + let (public_input, private_input) = prepare_inputs(); + let whir_config_builder = whir_config_builder(); + do_test_range_check( + 3716, + 20122, + &whir_config_builder, + &public_input, + &private_input, + ); +} + +#[test] +#[should_panic] +fn test_range_check_invalid_1() { + let (public_input, private_input) = prepare_inputs(); + let whir_config_builder = whir_config_builder(); + do_test_range_check( + 1, + 0, + &whir_config_builder, + &public_input, + &private_input, + ); +} diff --git a/crates/lean_vm/src/diagnostics/error.rs b/crates/lean_vm/src/diagnostics/error.rs index cc8ed616..9ce3b54f 100644 --- a/crates/lean_vm/src/diagnostics/error.rs +++ b/crates/lean_vm/src/diagnostics/error.rs @@ -1,3 +1,4 @@ +use std::collections::BTreeSet; use crate::core::F; use crate::diagnostics::profiler::MemoryProfile; use crate::execution::Memory; @@ -48,4 +49,5 @@ pub struct ExecutionResult { pub multilinear_evals: Vec, pub summary: String, pub memory_profile: Option, + pub range_check_cells_to_fill: Vec<(usize, usize)>, } diff --git a/crates/lean_vm/src/execution/runner.rs b/crates/lean_vm/src/execution/runner.rs index d19064f7..94c61511 100644 --- a/crates/lean_vm/src/execution/runner.rs +++ b/crates/lean_vm/src/execution/runner.rs @@ -214,6 +214,7 @@ fn execute_bytecode_helper( let mut counter_hint = 0; let mut cpu_cycles_before_new_line = 0; + let mut range_check_cells_to_fill = Vec::new(); while pc != bytecode.ending_pc { if pc >= bytecode.instructions.len() { @@ -264,6 +265,7 @@ fn execute_bytecode_helper( poseidon24_precomputed: poseidons_24_precomputed, n_poseidon16_precomputed_used: &mut n_poseidon16_precomputed_used, n_poseidon24_precomputed_used: &mut n_poseidon24_precomputed_used, + range_check_cells_to_fill: &mut range_check_cells_to_fill, }; instruction.execute_instruction(&mut instruction_ctx)?; } @@ -426,5 +428,6 @@ fn execute_bytecode_helper( multilinear_evals, summary, memory_profile: if profiling { Some(mem_profile) } else { None }, + range_check_cells_to_fill, }) } diff --git a/crates/lean_vm/src/isa/instruction.rs b/crates/lean_vm/src/isa/instruction.rs index 3b326e1d..2666ac6f 100644 --- a/crates/lean_vm/src/isa/instruction.rs +++ b/crates/lean_vm/src/isa/instruction.rs @@ -120,6 +120,7 @@ pub struct InstructionContext<'a> { pub poseidon24_precomputed: &'a [([F; 24], [F; 8])], pub n_poseidon16_precomputed_used: &'a mut usize, pub n_poseidon24_precomputed_used: &'a mut usize, + pub range_check_cells_to_fill: &'a mut Vec<(usize, usize)>, } impl Instruction { @@ -184,13 +185,18 @@ impl Instruction { let ptr = ctx.memory.get(*ctx.fp + shift_0)?; if *for_range_check { - // if for_range_check, ignore the UndefinedMemory error from get() - let value = ctx.memory.get(ptr.to_usize() + shift_1); + let ptr_usize = ptr.to_usize(); + let value = ctx.memory.get(ptr_usize + shift_1); if let Ok(value) = value { ctx.memory.set(memory_address_res, value)?; + } else { + // Ignore the UndefinedMemory error from get(). Also, indicate/"hint" + // to the prover that it needs to be filled later on with either 0 or + // some other value which a later instruction will set + ctx.range_check_cells_to_fill.push((memory_address_res, ptr_usize)); } } else { - // else, bubble it up + // For non-range check derefs, allow the error to bubble up let value = ctx.memory.get(ptr.to_usize() + shift_1)?; ctx.memory.set(memory_address_res, value)?; } @@ -212,6 +218,7 @@ impl Instruction { } => { let condition_value = condition.read_value(ctx.memory, *ctx.fp)?; assert!([F::ZERO, F::ONE].contains(&condition_value),); + if condition_value == F::ZERO { *ctx.pc += 1; } else { From ab5a4a3b90af30163073b4914d6c300966890cca Mon Sep 17 00:00:00 2001 From: Koh Wei Jie Date: Sun, 16 Nov 2025 16:49:19 -0800 Subject: [PATCH 10/10] cargo fmt --- crates/lean_prover/tests/test_range_check.rs | 10 ++-------- crates/lean_vm/src/diagnostics/error.rs | 2 +- crates/lean_vm/src/isa/instruction.rs | 3 ++- 3 files changed, 5 insertions(+), 10 deletions(-) diff --git a/crates/lean_prover/tests/test_range_check.rs b/crates/lean_prover/tests/test_range_check.rs index f5889743..8716e6f4 100644 --- a/crates/lean_prover/tests/test_range_check.rs +++ b/crates/lean_prover/tests/test_range_check.rs @@ -2,11 +2,11 @@ use lean_compiler::compile_program; use lean_prover::verify_execution::verify_execution; use lean_prover::{prove_execution::prove_execution, whir_config_builder}; use lean_vm::{DIMENSION, F, NONRESERVED_PROGRAM_INPUT_START}; +use multilinear_toolkit::prelude::*; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; use std::collections::BTreeSet; use whir_p3::WhirConfigBuilder; -use multilinear_toolkit::prelude::*; const NO_VEC_RUNTIME_MEMORY: usize = 1 << 20; @@ -156,11 +156,5 @@ fn test_range_check_valid_1() { fn test_range_check_invalid_1() { let (public_input, private_input) = prepare_inputs(); let whir_config_builder = whir_config_builder(); - do_test_range_check( - 1, - 0, - &whir_config_builder, - &public_input, - &private_input, - ); + do_test_range_check(1, 0, &whir_config_builder, &public_input, &private_input); } diff --git a/crates/lean_vm/src/diagnostics/error.rs b/crates/lean_vm/src/diagnostics/error.rs index 9ce3b54f..5519d53b 100644 --- a/crates/lean_vm/src/diagnostics/error.rs +++ b/crates/lean_vm/src/diagnostics/error.rs @@ -1,10 +1,10 @@ -use std::collections::BTreeSet; use crate::core::F; use crate::diagnostics::profiler::MemoryProfile; use crate::execution::Memory; use crate::witness::{ WitnessDotProduct, WitnessMultilinearEval, WitnessPoseidon16, WitnessPoseidon24, }; +use std::collections::BTreeSet; use thiserror::Error; #[derive(Debug, Clone, Error)] diff --git a/crates/lean_vm/src/isa/instruction.rs b/crates/lean_vm/src/isa/instruction.rs index 36825315..b0352b0a 100644 --- a/crates/lean_vm/src/isa/instruction.rs +++ b/crates/lean_vm/src/isa/instruction.rs @@ -192,7 +192,8 @@ impl Instruction { // Ignore the UndefinedMemory error from get(). Also, indicate/"hint" // to the prover that it needs to be filled later on with either 0 or // some other value which a later instruction will set - ctx.range_check_cells_to_fill.push((memory_address_res, ptr_usize)); + ctx.range_check_cells_to_fill + .push((memory_address_res, ptr_usize)); } } else { // For non-range check derefs, allow the error to bubble up