fix edge case: make the proof work even if there is no dot_product (precompile) calls during execution

Author: TomWambsgans

crates/lean_prover/src/common.rs

@@ -26,6 +26,7 @@ pub fn get_base_dims(
 ) -> Vec<ColDims<F>> {
     let n_poseidons_16 = n_poseidons_16.max(1 << LOG_MIN_POSEIDONS_16);
     let n_poseidons_24 = n_poseidons_24.max(1 << LOG_MIN_POSEIDONS_24);
+    let n_rows_table_dot_products = n_rows_table_dot_products.max(1 << LOG_MIN_DOT_PRODUCT_ROWS);
 
     let p16_default_cubes = default_cube_layers::<F, 16, N_COMMITED_CUBES_P16>(p16_gkr_layers);
     let p24_default_cubes = default_cube_layers::<F, 24, N_COMMITED_CUBES_P24>(p24_gkr_layers);

crates/lean_prover/src/prove_execution.rs

@@ -41,8 +41,8 @@ pub fn prove_execution(
         dot_products,
         multilinear_evals: vm_multilinear_evals,
         public_memory_size,
-        non_zero_memory_size,
-        memory, // padded with zeros to next power of two
+        mut non_zero_memory_size,
+        mut memory, // padded with zeros to next power of two
     } = info_span!("Witness generation").in_scope(|| {
         let mut execution_result = info_span!("Executing bytecode").in_scope(|| {
             execute_bytecode(
@@ -58,6 +58,10 @@ pub fn prove_execution(
             .in_scope(|| get_execution_trace(bytecode, execution_result))
     });
 
+    if memory.len() < 1 << MIN_LOG_MEMORY_SIZE {
+        memory.resize(1 << MIN_LOG_MEMORY_SIZE, F::ZERO);
+        non_zero_memory_size = 1 << MIN_LOG_MEMORY_SIZE;
+    }
     let public_memory = &memory[..public_memory_size];
     let private_memory = &memory[public_memory_size..non_zero_memory_size];
     let log_memory = log2_strict_usize(memory.len());
@@ -95,7 +99,8 @@ pub fn prove_execution(
 
     let dot_product_table = AirTable::<EF, _>::new(DotProductAir);
 
-    let (dot_product_columns, dot_product_padding_len) = build_dot_product_columns(&dot_products);
+    let (dot_product_columns, dot_product_padding_len) =
+        build_dot_product_columns(&dot_products, 1 << LOG_MIN_DOT_PRODUCT_ROWS);
 
     let dot_product_flags: Vec<PF<EF>> = field_slice_as_base(&dot_product_columns[0]).unwrap();
     let dot_product_lengths: Vec<PF<EF>> = field_slice_as_base(&dot_product_columns[1]).unwrap();

crates/lean_prover/src/verify_execution.rs

@@ -52,7 +52,10 @@ pub fn verify_execution(
 
     if log_n_cycles > 32
         || n_poseidons_16 > 1 << 32
-        || n_compressions_16 > n_poseidons_16.next_power_of_two()
+        || n_compressions_16
+            > n_poseidons_16
+                .next_power_of_two()
+                .max(1 << LOG_MIN_POSEIDONS_16)
         || n_poseidons_24 > 1 << 32
         || n_dot_products > 1 << 32
         || n_rows_table_dot_products > 1 << 32
@@ -75,9 +78,9 @@ pub fn verify_execution(
         return Err(ProofError::InvalidProof);
     }
 
-    let table_dot_products_log_n_rows = log2_ceil_usize(n_rows_table_dot_products);
-    let dot_product_padding_len =
-        n_rows_table_dot_products.next_power_of_two() - n_rows_table_dot_products;
+    let table_dot_products_log_n_rows =
+        log2_ceil_usize(n_rows_table_dot_products).max(LOG_MIN_DOT_PRODUCT_ROWS);
+    let dot_product_padding_len = (1 << table_dot_products_log_n_rows) - n_rows_table_dot_products;
 
     let mut vm_multilinear_evals = Vec::new();
     for _ in 0..n_vm_multilinear_evals {
@@ -204,7 +207,7 @@ pub fn verify_execution(
         p16_grand_product_evals_on_indexes_res,
     ] = verifier_state.next_extension_scalars_const().unwrap();
     let p16_grand_product_evals_on_compression = mle_of_zeros_then_ones(
-        n_poseidons_16.next_power_of_two() - n_compressions_16,
+        (1 << log_n_p16) - n_compressions_16,
         &grand_product_p16_statement.point,
     );
 

crates/lean_prover/tests/test_zkvm.rs

@@ -8,7 +8,7 @@ use p3_field::PrimeCharacteristicRing;
 const NO_VEC_RUNTIME_MEMORY: usize = 1 << 20;
 
 #[test]
-fn test_zk_vm() {
+fn test_zk_vm_all_precompiles() {
     let program_str = r#"
     const DIM = 5;
     const SECOND_POINT = 2;
@@ -59,9 +59,27 @@ fn test_zk_vm() {
 
         return;
     }
-   "#
-    .to_string();
+   "#;
 
+    test_zk_vm_helper(program_str);
+}
+
+#[test]
+fn test_zk_vm_no_precompiles() {
+    let program_str = r#"
+    
+    fn main() {
+        for i in 0..1000 {
+            assert i != 1000;
+        }
+        return;
+    }
+   "#;
+
+    test_zk_vm_helper(program_str);
+}
+
+fn test_zk_vm_helper(program_str: &str) {
     const SECOND_POINT: usize = 2;
     const SECOND_N_VARS: usize = 7;
 
@@ -82,7 +100,7 @@ fn test_zk_vm() {
         .collect::<Vec<_>>();
 
     // utils::init_tracing();
-    let bytecode = compile_program(program_str);
+    let bytecode = compile_program(program_str.to_string());
     let proof_data = prove_execution(
         &bytecode,
         (&public_input, &private_input),

crates/lean_prover/vm_air/src/dot_product_air.rs

@@ -96,7 +96,10 @@ impl<AB: AirBuilder> Air<AB> for DotProductAir {
     }
 }
 
-pub fn build_dot_product_columns(witness: &[WitnessDotProduct]) -> (Vec<Vec<EF>>, usize) {
+pub fn build_dot_product_columns(
+    witness: &[WitnessDotProduct],
+    min_n_rows: usize,
+) -> (Vec<Vec<EF>>, usize) {
     let (
         mut flag,
         mut len,
@@ -149,7 +152,7 @@ pub fn build_dot_product_columns(witness: &[WitnessDotProduct]) -> (Vec<Vec<EF>>
         res.extend(vec![dot_product.res; dot_product.len]);
     }
 
-    let padding_len = flag.len().next_power_of_two() - flag.len();
+    let padding_len = flag.len().next_power_of_two().max(min_n_rows) - flag.len();
     flag.extend(vec![EF::ONE; padding_len]);
     len.extend(vec![EF::ONE; padding_len]);
     index_a.extend(EF::zero_vec(padding_len));

crates/lean_prover/witness_generation/src/lib.rs

@@ -64,5 +64,7 @@ impl InAirColumnIndex for usize {
 
 // Zero padding will be added to each at least, if this minimum is not reached
 // (ensuring AIR / GKR work fine, with SIMD, without too much edge cases)
+// Long term, we should find a more elegant solution.
 pub const LOG_MIN_POSEIDONS_16: usize = 8;
 pub const LOG_MIN_POSEIDONS_24: usize = 8;
+pub const LOG_MIN_DOT_PRODUCT_ROWS: usize = 8;