implement some suggestions from @this-vishalsingh

Co-authored-by: thisvishalsingh <93567955+this-vishalsingh@users.noreply.github.com>

Author: TomWambsgans

crates/backend/fiat-shamir/src/lib.rs

@@ -20,3 +20,5 @@ pub(crate) use merkle_pruning::*;
 
 mod verifier;
 pub use verifier::*;
+
+const _: () = assert!(usize::BITS >= 32); // PoW grinding / Whir merkle index never exceeds 24 bits < 32

crates/lean_prover/src/prove_execution.rs

@@ -44,7 +44,11 @@ pub fn prove_execution(
     prover_state.observe_scalars(&poseidon16_compress_pair(&bytecode.hash, &SNARK_DOMAIN_SEP));
     prover_state.add_base_scalars(
         &[
-            vec![whir_config.starting_log_inv_rate, log2_strict_usize(memory.len())],
+            vec![
+                whir_config.starting_log_inv_rate,
+                log2_strict_usize(memory.len()),
+                public_input.len(),
+            ],
             traces.values().map(|t| t.log_n_rows).collect::<Vec<_>>(),
         ]
         .concat()

crates/lean_prover/src/verify_execution.rs

@@ -21,13 +21,17 @@ pub fn verify_execution(
     verifier_state.observe_scalars(public_input);
     verifier_state.observe_scalars(&poseidon16_compress_pair(&bytecode.hash, &SNARK_DOMAIN_SEP));
     let dims = verifier_state
-        .next_base_scalars_vec(2 + N_TABLES)?
+        .next_base_scalars_vec(3 + N_TABLES)?
         .into_iter()
         .map(|x| x.to_usize())
         .collect::<Vec<_>>();
     let log_inv_rate = dims[0];
     let log_memory = dims[1];
-    let table_n_vars: BTreeMap<Table, VarCount> = (0..N_TABLES).map(|i| (ALL_TABLES[i], dims[i + 2])).collect();
+    let public_input_len = dims[2]; // enforce the exact length of the public input to pass through Fiat Shamir (otherwise we could have 2 public inputs, only differing by a few (<8) zeros in the end, leading to the same fiat shamir state: tipically giving the advseary 2 or 3 bits of advantage in the subsequent part where the public input is evaluated as a multilinear polynomial)
+    if public_input_len != public_input.len() {
+        return Err(ProofError::InvalidProof);
+    }
+    let table_n_vars: BTreeMap<Table, VarCount> = (0..N_TABLES).map(|i| (ALL_TABLES[i], dims[i + 3])).collect();
     if !(MIN_WHIR_LOG_INV_RATE..=MAX_WHIR_LOG_INV_RATE).contains(&log_inv_rate) {
         return Err(ProofError::InvalidProof);
     }

crates/rec_aggregation/recursion.py

@@ -52,11 +52,14 @@ def recursion(inner_public_memory, proof_transcript, bytecode_value_hint):
     # table dims
     debug_assert(N_TABLES + 1 < DIGEST_LEN)
     fs, dims = fs_receive_chunks(fs, 1)
-    for i in unroll(N_TABLES + 2, 8):
+    for i in unroll(N_TABLES + 3, 8):
         assert dims[i] == 0
     whir_log_inv_rate = dims[0]
     log_memory = dims[1]
-    table_log_heights = dims + 2
+    public_input_len = dims[2]
+    table_log_heights = dims + 3
+
+    assert public_input_len == PUB_INPUT_SIZE
 
     assert MIN_WHIR_LOG_INV_RATE <= whir_log_inv_rate
     assert whir_log_inv_rate <= MAX_WHIR_LOG_INV_RATE