Enables optional verification of Keccak 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,
 };
@@ -1916,30 +1916,59 @@ where
         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())?;
         }
@@ -1958,8 +1987,7 @@ where
             anyhow::Error::msg("Invalid conversion when setting public values targets.")
         })?;
 
-        // TODO(sdeng): Set to false when this segment contains no Keccak operations.
-        root_inputs.set_bool_target(self.root.use_keccak_tables, true);
+        root_inputs.set_bool_target(self.root.use_keccak_tables, all_proof.use_keccak_tables);
 
         let root_proof = self.root.circuit.prove(root_inputs)?;
 
@@ -2075,16 +2103,36 @@ where
 
         for table in 0..NUM_TABLES {
             let (table_circuit, index_verifier_data) = &table_circuits[table];
+            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 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);
+                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())?;
         }
@@ -2103,8 +2151,7 @@ where
             anyhow::Error::msg("Invalid conversion when setting public values targets.")
         })?;
 
-        // TODO(sdeng): Set to false when this segment contains no Keccak operations.
-        root_inputs.set_bool_target(self.root.use_keccak_tables, true);
+        root_inputs.set_bool_target(self.root.use_keccak_tables, all_proof.use_keccak_tables);
 
         let root_proof = self.root.circuit.prove(root_inputs)?;
 
@@ -3092,7 +3139,6 @@ mod tests {
     type C = PoseidonGoldilocksConfig;
 
     #[test]
-    #[ignore]
     fn test_segment_proof_generation_without_keccak() -> anyhow::Result<()> {
         init_logger();
 

evm_arithmetization/src/generation/mod.rs

@@ -486,15 +486,19 @@ fn get_all_memory_address_and_values(memory_before: &MemoryState) -> Vec<(Memory
     res
 }
 
-type TablesWithPVsAndFinalMem<F> = ([Vec<PolynomialValues<F>>; NUM_TABLES], PublicValues<F>);
+pub struct TablesWithPVs<F: RichField> {
+    pub tables: [Vec<PolynomialValues<F>>; NUM_TABLES],
+    pub use_keccak_tables: bool,
+    pub public_values: PublicValues<F>,
+}
 
 pub fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
     all_stark: &AllStark<F, D>,
     inputs: &TrimmedGenerationInputs<F>,
     config: &StarkConfig,
     segment_data: &mut GenerationSegmentData,
     timing: &mut TimingTree,
-) -> anyhow::Result<TablesWithPVsAndFinalMem<F>> {
+) -> anyhow::Result<TablesWithPVs<F>> {
     let mut state = GenerationState::<F>::new_with_segment_data(inputs, segment_data)
         .map_err(|err| anyhow!("Failed to parse all the initial prover inputs: {:?}", err))?;
 
@@ -581,6 +585,9 @@ pub fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
         mem_after: MemCap::default(),
     };
 
+    let use_keccak_tables =
+        !state.traces.keccak_inputs.is_empty() || !state.traces.keccak_sponge_ops.is_empty();
+
     let tables = timed!(
         timing,
         "convert trace data to tables",
@@ -593,7 +600,12 @@ pub fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
             timing
         )
     );
-    Ok((tables, public_values))
+
+    Ok(TablesWithPVs {
+        tables,
+        use_keccak_tables,
+        public_values,
+    })
 }
 
 fn simulate_cpu<F: RichField>(

evm_arithmetization/src/get_challenges.rs

@@ -6,6 +6,7 @@ use plonky2::plonk::config::{AlgebraicHasher, GenericConfig};
 use starky::config::StarkConfig;
 use starky::lookup::get_grand_product_challenge_set;
 
+use crate::all_stark::KECCAK_TABLES_INDICES;
 use crate::proof::*;
 use crate::util::{h256_limbs, u256_limbs, u256_to_u32, u256_to_u64};
 use crate::witness::errors::ProgramError;
@@ -257,8 +258,20 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize> A
 
         let stark_proofs = &self.multi_proof.stark_proofs;
 
-        for proof in stark_proofs {
-            challenger.observe_cap(&proof.proof.trace_cap);
+        for (i, proof) in stark_proofs.iter().enumerate() {
+            if KECCAK_TABLES_INDICES.contains(&i) && !self.use_keccak_tables {
+                // Observe zero merkle caps when skipping Keccak tables.
+                let zero_merkle_cap = proof
+                    .proof
+                    .trace_cap
+                    .flatten()
+                    .iter()
+                    .map(|_| F::ZERO)
+                    .collect::<Vec<F>>();
+                challenger.observe_elements(&zero_merkle_cap);
+            } else {
+                challenger.observe_cap(&proof.proof.trace_cap);
+            }
         }
 
         observe_public_values::<F, C, D>(&mut challenger, &self.public_values)?;
@@ -268,13 +281,17 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize> A
 
         Ok(AllProofChallenges {
             stark_challenges: core::array::from_fn(|i| {
-                challenger.compact();
-                stark_proofs[i].proof.get_challenges(
-                    &mut challenger,
-                    Some(&ctl_challenges),
-                    true,
-                    config,
-                )
+                if KECCAK_TABLES_INDICES.contains(&i) && !self.use_keccak_tables {
+                    None
+                } else {
+                    challenger.compact();
+                    Some(stark_proofs[i].proof.get_challenges(
+                        &mut challenger,
+                        Some(&ctl_challenges),
+                        true,
+                        config,
+                    ))
+                }
             }),
             ctl_challenges,
         })

evm_arithmetization/src/proof.rs

@@ -31,6 +31,10 @@ pub struct AllProof<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, co
     pub multi_proof: MultiProof<F, C, D, NUM_TABLES>,
     /// Public memory values used for the recursive proofs.
     pub public_values: PublicValues<F>,
+    /// A flag indicating whether the Keccak and KeccakSponge tables contain
+    /// only padding values (i.e., no meaningful data). This is set to false
+    /// when no actual Keccak operations were performed.
+    pub use_keccak_tables: bool,
 }
 
 impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize> AllProof<F, C, D> {
@@ -43,7 +47,7 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize> A
 /// Randomness for all STARKs.
 pub(crate) struct AllProofChallenges<F: RichField + Extendable<D>, const D: usize> {
     /// Randomness used in each STARK proof.
-    pub stark_challenges: [StarkProofChallenges<F, D>; NUM_TABLES],
+    pub stark_challenges: [Option<StarkProofChallenges<F, D>>; NUM_TABLES],
     /// Randomness used for cross-table lookups. It is shared by all STARKs.
     pub ctl_challenges: GrandProductChallengeSet<F>,
 }

evm_arithmetization/src/prover.rs

@@ -20,7 +20,7 @@ use starky::proof::{MultiProof, StarkProofWithMetadata};
 use starky::prover::prove_with_commitment;
 use starky::stark::Stark;
 
-use crate::all_stark::{AllStark, Table, NUM_TABLES};
+use crate::all_stark::{AllStark, Table, KECCAK_TABLES_INDICES, NUM_TABLES};
 use crate::cpu::kernel::aggregator::KERNEL;
 use crate::generation::segments::GenerationSegmentData;
 use crate::generation::{generate_traces, GenerationInputs, TrimmedGenerationInputs};
@@ -47,7 +47,7 @@ where
 
     timed!(timing, "build kernel", Lazy::force(&KERNEL));
 
-    let (traces, mut public_values) = timed!(
+    let mut tables_with_pvs = timed!(
         timing,
         "generate all traces",
         generate_traces(all_stark, &inputs, config, segment_data, timing)?
@@ -58,8 +58,9 @@ where
     let proof = prove_with_traces(
         all_stark,
         config,
-        traces,
-        &mut public_values,
+        tables_with_pvs.tables,
+        tables_with_pvs.use_keccak_tables,
+        &mut tables_with_pvs.public_values,
         timing,
         abort_signal,
     )?;
@@ -72,6 +73,7 @@ pub(crate) fn prove_with_traces<F, C, const D: usize>(
     all_stark: &AllStark<F, D>,
     config: &StarkConfig,
     trace_poly_values: [Vec<PolynomialValues<F>>; NUM_TABLES],
+    use_keccak_tables: bool,
     public_values: &mut PublicValues<F>,
     timing: &mut TimingTree,
     abort_signal: Option<Arc<AtomicBool>>,
@@ -114,8 +116,14 @@ where
         .map(|c| c.merkle_tree.cap.clone())
         .collect::<Vec<_>>();
     let mut challenger = Challenger::<F, C::Hasher>::new();
-    for cap in &trace_caps {
-        challenger.observe_cap(cap);
+    for (i, cap) in trace_caps.iter().enumerate() {
+        if KECCAK_TABLES_INDICES.contains(&i) && !use_keccak_tables {
+            // Observe zero merkle caps when skipping Keccak tables.
+            let zero_merkle_cap = cap.flatten().iter().map(|_| F::ZERO).collect::<Vec<F>>();
+            challenger.observe_elements(&zero_merkle_cap);
+        } else {
+            challenger.observe_cap(cap);
+        }
     }
 
     observe_public_values::<F, C, D>(&mut challenger, public_values)
@@ -143,6 +151,7 @@ where
             config,
             &trace_poly_values,
             trace_commitments,
+            use_keccak_tables,
             ctl_data_per_table,
             &mut challenger,
             &ctl_challenges,
@@ -206,6 +215,7 @@ where
             ctl_challenges,
         },
         public_values: public_values.clone(),
+        use_keccak_tables,
     })
 }
 
@@ -229,6 +239,7 @@ fn prove_with_commitments<F, C, const D: usize>(
     config: &StarkConfig,
     trace_poly_values: &[Vec<PolynomialValues<F>>; NUM_TABLES],
     trace_commitments: Vec<PolynomialBatch<F, C, D>>,
+    use_keccak_tables: bool,
     ctl_data_per_table: [CtlData<F>; NUM_TABLES],
     challenger: &mut Challenger<F, C::Hasher>,
     ctl_challenges: &GrandProductChallengeSet<F>,
@@ -262,8 +273,16 @@ where
     let (arithmetic_proof, _) = prove_table!(arithmetic_stark, Table::Arithmetic);
     let (byte_packing_proof, _) = prove_table!(byte_packing_stark, Table::BytePacking);
     let (cpu_proof, _) = prove_table!(cpu_stark, Table::Cpu);
+    let challenger_after_cpu = challenger.clone();
+    // TODO(sdeng): Keccak proofs are still required for CTLs, etc. Refactor the
+    // code and remove the unnecessary parts.
     let (keccak_proof, _) = prove_table!(keccak_stark, Table::Keccak);
     let (keccak_sponge_proof, _) = prove_table!(keccak_sponge_stark, Table::KeccakSponge);
+    if !use_keccak_tables {
+        // We need to connect the challenger state of Logic and CPU tables when the
+        // Keccak tables are not in use.
+        *challenger = challenger_after_cpu;
+    }
     let (logic_proof, _) = prove_table!(logic_stark, Table::Logic);
     let (memory_proof, _) = prove_table!(memory_stark, Table::Memory);
     let (mem_before_proof, mem_before_cap) = prove_table!(mem_before_stark, Table::MemBefore);