fix edge case: make the proof work even if there is no Poseidon calls during execution

Author: TomWambsgans

Cargo.lock

@@ -17,14 +17,16 @@ pub fn get_base_dims(
     log_public_memory: usize,
     private_memory_len: usize,
     bytecode_ending_pc: usize,
-    poseidon_counts: (usize, usize),
+    (n_poseidons_16, n_poseidons_24): (usize, usize),
     n_rows_table_dot_products: usize,
     (p16_gkr_layers, p24_gkr_layers): (
         &PoseidonGKRLayers<16, N_COMMITED_CUBES_P16>,
         &PoseidonGKRLayers<24, N_COMMITED_CUBES_P24>,
     ),
 ) -> Vec<ColDims<F>> {
-    let (n_poseidons_16, n_poseidons_24) = poseidon_counts;
+    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 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/common.rs

@@ -67,8 +67,8 @@ pub fn prove_execution(
     let log_n_cycles = log2_strict_usize(n_cycles);
     assert!(full_trace.iter().all(|col| col.len() == 1 << log_n_cycles));
 
-    let log_n_p16 = log2_ceil_usize(n_poseidons_16);
-    let log_n_p24 = log2_ceil_usize(n_poseidons_24);
+    let log_n_p16 = log2_ceil_usize(n_poseidons_16).max(LOG_MIN_POSEIDONS_16);
+    let log_n_p24 = log2_ceil_usize(n_poseidons_24).max(LOG_MIN_POSEIDONS_24);
 
     precompute_dft_twiddles::<F>(1 << 24);
 
@@ -334,7 +334,7 @@ pub fn prove_execution(
                 &WitnessPoseidon16::default_addresses_field_repr(POSEIDON_16_NULL_HASH_PTR),
                 fingerprint_challenge,
             ))
-        .exp_u64((n_poseidons_16.next_power_of_two() - n_poseidons_16) as u64);
+        .exp_u64(((1 << log_n_p16) - n_poseidons_16) as u64);
 
     let corrected_prod_p24 = grand_product_p24_res
         / (grand_product_challenge_global
@@ -343,7 +343,7 @@ pub fn prove_execution(
                 &WitnessPoseidon24::default_addresses_field_repr(POSEIDON_24_NULL_HASH_PTR),
                 fingerprint_challenge,
             ))
-        .exp_u64((n_poseidons_24.next_power_of_two() - n_poseidons_24) as u64);
+        .exp_u64(((1 << log_n_p24) - n_poseidons_24) as u64);
 
     let corrected_dot_product = grand_product_dot_product_res
         / ((grand_product_challenge_global

crates/lean_prover/src/prove_execution.rs

@@ -68,8 +68,8 @@ pub fn verify_execution(
 
     let log_public_memory = log2_strict_usize(public_memory.len());
     let log_memory = log2_ceil_usize(public_memory.len() + private_memory_len);
-    let log_n_p16 = log2_ceil_usize(n_poseidons_16);
-    let log_n_p24 = log2_ceil_usize(n_poseidons_24);
+    let log_n_p16 = log2_ceil_usize(n_poseidons_16).max(LOG_MIN_POSEIDONS_16);
+    let log_n_p24 = log2_ceil_usize(n_poseidons_24).max(LOG_MIN_POSEIDONS_24);
 
     if !(MIN_LOG_MEMORY_SIZE..=MAX_LOG_MEMORY_SIZE).contains(&log_memory) {
         return Err(ProofError::InvalidProof);
@@ -159,7 +159,7 @@ pub fn verify_execution(
                 &WitnessPoseidon16::default_addresses_field_repr(POSEIDON_16_NULL_HASH_PTR),
                 fingerprint_challenge,
             ))
-        .exp_u64((n_poseidons_16.next_power_of_two() - n_poseidons_16) as u64);
+        .exp_u64(((1 << log_n_p16) - n_poseidons_16) as u64);
 
     let corrected_prod_p24 = grand_product_p24_res
         / (grand_product_challenge_global
@@ -168,7 +168,7 @@ pub fn verify_execution(
                 &WitnessPoseidon24::default_addresses_field_repr(POSEIDON_24_NULL_HASH_PTR),
                 fingerprint_challenge,
             ))
-        .exp_u64((n_poseidons_24.next_power_of_two() - n_poseidons_24) as u64);
+        .exp_u64(((1 << log_n_p24) - n_poseidons_24) as u64);
 
     let corrected_dot_product = grand_product_dot_product_res
         / ((grand_product_challenge_global

crates/lean_prover/src/verify_execution.rs

@@ -4,7 +4,7 @@ use crate::instruction_encoder::field_representation;
 use crate::{
     COL_INDEX_FP, COL_INDEX_MEM_ADDRESS_A, COL_INDEX_MEM_ADDRESS_B, COL_INDEX_MEM_ADDRESS_C,
     COL_INDEX_MEM_VALUE_A, COL_INDEX_MEM_VALUE_B, COL_INDEX_MEM_VALUE_C, COL_INDEX_PC,
-    N_TOTAL_COLUMNS,
+    LOG_MIN_POSEIDONS_16, LOG_MIN_POSEIDONS_24, N_TOTAL_COLUMNS,
 };
 use lean_vm::*;
 use p3_field::Field;
@@ -108,14 +108,18 @@ pub fn get_execution_trace(
         ..
     } = execution_result;
 
-    poseidons_16.extend(vec![
-        WitnessPoseidon16::poseidon_of_zero();
-        n_poseidons_16.next_power_of_two() - n_poseidons_16
-    ]);
-    poseidons_24.extend(vec![
-        WitnessPoseidon24::poseidon_of_zero();
-        n_poseidons_24.next_power_of_two() - n_poseidons_24
-    ]);
+    poseidons_16.resize(
+        n_poseidons_16
+            .next_power_of_two()
+            .max(1 << LOG_MIN_POSEIDONS_16),
+        WitnessPoseidon16::poseidon_of_zero(),
+    );
+    poseidons_24.resize(
+        n_poseidons_24
+            .next_power_of_two()
+            .max(1 << LOG_MIN_POSEIDONS_24),
+        WitnessPoseidon24::poseidon_of_zero(),
+    );
 
     let n_compressions_16;
     (poseidons_16, n_compressions_16) = put_poseidon16_compressions_at_the_end(&poseidons_16); // TODO avoid reallocation

crates/lean_prover/witness_generation/src/execution_trace.rs

@@ -61,3 +61,8 @@ 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)
+pub const LOG_MIN_POSEIDONS_16: usize = 8;
+pub const LOG_MIN_POSEIDONS_24: usize = 8;

crates/lean_prover/witness_generation/src/lib.rs

@@ -37,7 +37,7 @@ where
 {
     assert!(log2_strict_usize(final_layer.len()) >= 1);
     if final_layer.len() == 2 {
-        prover_state.add_extension_scalars(&final_layer);
+        prover_state.add_extension_scalars(final_layer);
         let product = final_layer[0] * final_layer[1];
         let point = MultilinearPoint(vec![prover_state.sample()]);
         let claim = Evaluation {
@@ -71,7 +71,7 @@ where
     };
     assert_eq!(last_layer.len(), 2);
     let product = last_layer[0] * last_layer[1];
-    prover_state.add_extension_scalars(&last_layer);
+    prover_state.add_extension_scalars(last_layer);
 
     let point = MultilinearPoint(vec![prover_state.sample()]);
     let mut claim = Evaluation {