Enables optional verification of Keccak tables

evm_arithmetization/src/all_stark.rs

@@ -126,6 +126,10 @@ pub const NUM_TABLES: usize = if cfg!(feature = "cdk_erigon") {
     Table::MemAfter as usize + 1
 };
 
+/// Indices of Keccak Tables
+pub const KECCAK_TABLES_INDICES: [usize; 2] =
+    [Table::Keccak as usize, Table::KeccakSponge as usize];
+
 impl Table {
     /// Returns all STARK table indices.
     pub(crate) const fn all() -> [Self; NUM_TABLES] {

evm_arithmetization/src/fixed_recursive_verifier.rs

@@ -23,7 +23,7 @@ use plonky2::plonk::circuit_data::{
 use plonky2::plonk::config::{AlgebraicHasher, GenericConfig, GenericHashOut};
 use plonky2::plonk::proof::{ProofWithPublicInputs, ProofWithPublicInputsTarget};
 use plonky2::recursion::cyclic_recursion::check_cyclic_proof_verifier_data;
-use plonky2::recursion::dummy_circuit::cyclic_base_proof;
+use plonky2::recursion::dummy_circuit::{cyclic_base_proof, dummy_circuit, dummy_proof};
 use plonky2::util::serialization::{
     Buffer, GateSerializer, IoResult, Read, WitnessGeneratorSerializer, Write,
 };
@@ -36,7 +36,9 @@ use starky::lookup::{get_grand_product_challenge_set_target, GrandProductChallen
 use starky::proof::StarkProofWithMetadata;
 use starky::stark::Stark;
 
-use crate::all_stark::{all_cross_table_lookups, AllStark, Table, NUM_TABLES};
+use crate::all_stark::{
+    all_cross_table_lookups, AllStark, Table, KECCAK_TABLES_INDICES, NUM_TABLES,
+};
 use crate::cpu::kernel::aggregator::KERNEL;
 use crate::generation::segments::{GenerationSegmentData, SegmentDataIterator};
 use crate::generation::{GenerationInputs, TrimmedGenerationInputs};
@@ -151,6 +153,8 @@ where
     /// for EVM root proofs; the circuit has them just to match the
     /// structure of aggregation proofs.
     cyclic_vk: VerifierCircuitTarget,
+    /// We can skip verifying Keccak tables when they are not in use.
+    use_keccak_tables: BoolTarget,
 }
 
 impl<F, C, const D: usize> RootCircuitData<F, C, D>
@@ -173,6 +177,7 @@ where
         }
         self.public_values.to_buffer(buffer)?;
         buffer.write_target_verifier_circuit(&self.cyclic_vk)?;
+        buffer.write_target_bool(self.use_keccak_tables)?;
         Ok(())
     }
 
@@ -192,13 +197,15 @@ where
         }
         let public_values = PublicValuesTarget::from_buffer(buffer)?;
         let cyclic_vk = buffer.read_target_verifier_circuit()?;
+        let use_keccak_tables = buffer.read_target_bool()?;
 
         Ok(Self {
             circuit,
             proof_with_pis: proof_with_pis.try_into().unwrap(),
             index_verifier_data: index_verifier_data.try_into().unwrap(),
             public_values,
             cyclic_vk,
+            use_keccak_tables,
         })
     }
 }
@@ -807,6 +814,8 @@ where
 
         let mut builder = CircuitBuilder::new(CircuitConfig::standard_recursion_config());
 
+        let use_keccak_tables = builder.add_virtual_bool_target_safe();
+        let skip_keccak_tables = builder.not(use_keccak_tables);
         let public_values = add_virtual_public_values_public_input(&mut builder);
 
         let recursive_proofs =
@@ -826,6 +835,16 @@ where
             }
         }
 
+        // Ensures that the trace cap is set to 0 when skipping Keccak tables.
+        for i in KECCAK_TABLES_INDICES {
+            for h in &pis[i].trace_cap {
+                for t in h {
+                    let trace_cap_check = builder.mul(skip_keccak_tables.target, *t);
+                    builder.assert_zero(trace_cap_check);
+                }
+            }
+        }
+
         observe_public_values_target::<F, C, D>(&mut challenger, &public_values);
 
         let ctl_challenges = get_grand_product_challenge_set_target(
@@ -834,16 +853,31 @@ where
             stark_config.num_challenges,
         );
         // Check that the correct CTL challenges are used in every proof.
-        for pi in &pis {
-            for i in 0..stark_config.num_challenges {
-                builder.connect(
-                    ctl_challenges.challenges[i].beta,
-                    pi.ctl_challenges.challenges[i].beta,
-                );
-                builder.connect(
-                    ctl_challenges.challenges[i].gamma,
-                    pi.ctl_challenges.challenges[i].gamma,
-                );
+        for (i, pi) in pis.iter().enumerate() {
+            for j in 0..stark_config.num_challenges {
+                if KECCAK_TABLES_INDICES.contains(&i) {
+                    // Ensures that the correct CTL challenges are used in Keccak tables when
+                    // `enable_keccak_tables` is true.
+                    builder.conditional_assert_eq(
+                        use_keccak_tables.target,
+                        ctl_challenges.challenges[j].beta,
+                        pi.ctl_challenges.challenges[j].beta,
+                    );
+                    builder.conditional_assert_eq(
+                        use_keccak_tables.target,
+                        ctl_challenges.challenges[j].gamma,
+                        pi.ctl_challenges.challenges[j].gamma,
+                    );
+                } else {
+                    builder.connect(
+                        ctl_challenges.challenges[j].beta,
+                        pi.ctl_challenges.challenges[j].beta,
+                    );
+                    builder.connect(
+                        ctl_challenges.challenges[j].gamma,
+                        pi.ctl_challenges.challenges[j].gamma,
+                    );
+                }
             }
         }
 
@@ -852,12 +886,28 @@ where
             builder.connect(before, s);
         }
         // Check that the challenger state is consistent between proofs.
+        let mut prev_state = pis[0].challenger_state_after.as_ref().to_vec();
+        let state_len = prev_state.len();
         for i in 1..NUM_TABLES {
-            for (&before, &after) in zip_eq(
-                pis[i].challenger_state_before.as_ref(),
-                pis[i - 1].challenger_state_after.as_ref(),
-            ) {
-                builder.connect(before, after);
+            let current_state_before = pis[i].challenger_state_before.as_ref();
+            let current_state_after = pis[i].challenger_state_after.as_ref();
+            for j in 0..state_len {
+                if KECCAK_TABLES_INDICES.contains(&i) {
+                    // Ensure the challenger state:
+                    // 1) prev == current_before when using Keccak
+                    builder.conditional_assert_eq(
+                        use_keccak_tables.target,
+                        prev_state[j],
+                        current_state_before[j],
+                    );
+                    // 2) Update prev <- current_after when using Keccak
+                    // 3) Keep prev <- prev when skipping Keccak
+                    prev_state[j] =
+                        builder.select(use_keccak_tables, current_state_after[j], prev_state[j]);
+                } else {
+                    builder.connect(prev_state[j], current_state_before[j]);
+                    prev_state[j] = current_state_after[j];
+                }
             }
         }
 
@@ -877,6 +927,15 @@ where
             })
             .collect_vec();
 
+        // Ensure that when Keccak tables are skipped, the Keccak tables' ctl_zs_first
+        // are all zeros.
+        for &i in KECCAK_TABLES_INDICES.iter() {
+            for &t in pis[i].ctl_zs_first.iter() {
+                let ctl_check = builder.mul(skip_keccak_tables.target, t);
+                builder.assert_zero(ctl_check);
+            }
+        }
+
         // Verify the CTL checks.
         verify_cross_table_lookups_circuit::<F, D, NUM_TABLES>(
             &mut builder,
@@ -906,11 +965,22 @@ where
             let inner_verifier_data =
                 builder.random_access_verifier_data(index_verifier_data[i], possible_vks);
 
-            builder.verify_proof::<C>(
-                &recursive_proofs[i],
-                &inner_verifier_data,
-                inner_common_data[i],
-            );
+            if KECCAK_TABLES_INDICES.contains(&i) {
+                builder
+                    .conditionally_verify_proof_or_dummy::<C>(
+                        use_keccak_tables,
+                        &recursive_proofs[i],
+                        &inner_verifier_data,
+                        inner_common_data[i],
+                    )
+                    .expect("Unable conditionally verify Keccak proofs in the root circuit");
+            } else {
+                builder.verify_proof::<C>(
+                    &recursive_proofs[i],
+                    &inner_verifier_data,
+                    inner_common_data[i],
+                );
+            }
         }
 
         let merkle_before =
@@ -941,6 +1011,7 @@ where
             index_verifier_data,
             public_values,
             cyclic_vk,
+            use_keccak_tables,
         }
     }
 
@@ -1841,33 +1912,63 @@ where
             timing,
             abort_signal.clone(),
         )?;
+
         let mut root_inputs = PartialWitness::new();
 
         for table in 0..NUM_TABLES {
-            let stark_proof = &all_proof.multi_proof.stark_proofs[table];
-            let original_degree_bits = stark_proof.proof.recover_degree_bits(config);
             let table_circuits = &self.by_table[table];
-            let shrunk_proof = table_circuits
-                .by_stark_size
-                .get(&original_degree_bits)
-                .ok_or_else(|| {
-                    anyhow!(format!(
-                        "Missing preprocessed circuits for {:?} table with size {}.",
-                        Table::all()[table],
-                        original_degree_bits,
-                    ))
-                })?
-                .shrink(stark_proof, &all_proof.multi_proof.ctl_challenges)?;
-            let index_verifier_data = table_circuits
-                .by_stark_size
-                .keys()
-                .position(|&size| size == original_degree_bits)
-                .unwrap();
-            root_inputs.set_target(
-                self.root.index_verifier_data[table],
-                F::from_canonical_usize(index_verifier_data),
-            );
-            root_inputs.set_proof_with_pis_target(&self.root.proof_with_pis[table], &shrunk_proof);
+            if KECCAK_TABLES_INDICES.contains(&table) && !all_proof.use_keccak_tables {
+                // generate and set a dummy `index_verifier_data` and `proof_with_pis`
+                let index_verifier_data = table_circuits
+                    .by_stark_size
+                    .keys()
+                    .min()
+                    .expect("No valid size in shrinking circuits");
+                root_inputs.set_target(
+                    self.root.index_verifier_data[table],
+                    F::from_canonical_usize(*index_verifier_data),
+                );
+                let common_data = &table_circuits
+                    .by_stark_size
+                    .get(index_verifier_data)
+                    .expect("No valid size in shrinking circuits")
+                    .shrinking_wrappers
+                    .last()
+                    .expect("No shrinking circuits")
+                    .circuit
+                    .common;
+                let dummy_circuit: CircuitData<F, C, D> = dummy_circuit(common_data);
+                let dummy_pis = HashMap::new();
+                let dummy_proof = dummy_proof(&dummy_circuit, dummy_pis)
+                    .expect("Unable to generate dummy proofs");
+                root_inputs
+                    .set_proof_with_pis_target(&self.root.proof_with_pis[table], &dummy_proof);
+            } else {
+                let stark_proof = &all_proof.multi_proof.stark_proofs[table];
+                let original_degree_bits = stark_proof.proof.recover_degree_bits(config);
+                let shrunk_proof = table_circuits
+                    .by_stark_size
+                    .get(&original_degree_bits)
+                    .ok_or_else(|| {
+                        anyhow!(format!(
+                            "Missing preprocessed circuits for {:?} table with size {}.",
+                            Table::all()[table],
+                            original_degree_bits,
+                        ))
+                    })?
+                    .shrink(stark_proof, &all_proof.multi_proof.ctl_challenges)?;
+                let index_verifier_data = table_circuits
+                    .by_stark_size
+                    .keys()
+                    .position(|&size| size == original_degree_bits)
+                    .unwrap();
+                root_inputs.set_target(
+                    self.root.index_verifier_data[table],
+                    F::from_canonical_usize(index_verifier_data),
+                );
+                root_inputs
+                    .set_proof_with_pis_target(&self.root.proof_with_pis[table], &shrunk_proof);
+            }
 
             check_abort_signal(abort_signal.clone())?;
         }
@@ -1886,6 +1987,8 @@ where
             anyhow::Error::msg("Invalid conversion when setting public values targets.")
         })?;
 
+        root_inputs.set_bool_target(self.root.use_keccak_tables, all_proof.use_keccak_tables);
+
         let root_proof = self.root.circuit.prove(root_inputs)?;
 
         Ok(ProverOutputData {
@@ -2000,16 +2103,36 @@ where
 
         for table in 0..NUM_TABLES {
             let (table_circuit, index_verifier_data) = &table_circuits[table];
-
-            let stark_proof = &all_proof.multi_proof.stark_proofs[table];
-
-            let shrunk_proof =
-                table_circuit.shrink(stark_proof, &all_proof.multi_proof.ctl_challenges)?;
-            root_inputs.set_target(
-                self.root.index_verifier_data[table],
-                F::from_canonical_u8(*index_verifier_data),
-            );
-            root_inputs.set_proof_with_pis_target(&self.root.proof_with_pis[table], &shrunk_proof);
+            if KECCAK_TABLES_INDICES.contains(&table) && !all_proof.use_keccak_tables {
+                root_inputs.set_target(
+                    self.root.index_verifier_data[table],
+                    F::from_canonical_u8(*index_verifier_data),
+                );
+                // generate and set a dummy `proof_with_pis`
+                let common_date = &table_circuit
+                    .shrinking_wrappers
+                    .last()
+                    .expect("No shrinking circuits")
+                    .circuit
+                    .common;
+                let dummy_circuit: CircuitData<F, C, D> = dummy_circuit(common_date);
+                let dummy_pis = HashMap::new();
+                let dummy_proof = dummy_proof(&dummy_circuit, dummy_pis)
+                    .expect("Unable to generate dummy proofs");
+                root_inputs
+                    .set_proof_with_pis_target(&self.root.proof_with_pis[table], &dummy_proof);
+            } else {
+                let stark_proof = &all_proof.multi_proof.stark_proofs[table];
+
+                let shrunk_proof =
+                    table_circuit.shrink(stark_proof, &all_proof.multi_proof.ctl_challenges)?;
+                root_inputs.set_target(
+                    self.root.index_verifier_data[table],
+                    F::from_canonical_u8(*index_verifier_data),
+                );
+                root_inputs
+                    .set_proof_with_pis_target(&self.root.proof_with_pis[table], &shrunk_proof);
+            }
 
             check_abort_signal(abort_signal.clone())?;
         }
@@ -2028,6 +2151,8 @@ where
             anyhow::Error::msg("Invalid conversion when setting public values targets.")
         })?;
 
+        root_inputs.set_bool_target(self.root.use_keccak_tables, all_proof.use_keccak_tables);
+
         let root_proof = self.root.circuit.prove(root_inputs)?;
 
         Ok(ProofWithPublicValues {
@@ -3014,7 +3139,6 @@ mod tests {
     type C = PoseidonGoldilocksConfig;
 
     #[test]
-    #[ignore]
     fn test_segment_proof_generation_without_keccak() -> anyhow::Result<()> {
         init_logger();