Make evm structs more generic

evm/src/all_stark.rs

@@ -16,10 +16,21 @@ pub struct AllStark<F: RichField + Extendable<D>, const D: usize> {
 
 impl<F: RichField + Extendable<D>, const D: usize> AllStark<F, D> {
     pub(crate) fn nums_permutation_zs(&self, config: &StarkConfig) -> Vec<usize> {
-        vec![
+        let ans = vec![
             self.cpu_stark.num_permutation_batches(config),
             self.keccak_stark.num_permutation_batches(config),
-        ]
+        ];
+        debug_assert_eq!(ans.len(), Table::num_tables());
+        ans
+    }
+
+    pub(crate) fn permutation_batch_sizes(&self) -> Vec<usize> {
+        let ans = vec![
+            self.cpu_stark.permutation_batch_size(),
+            self.keccak_stark.permutation_batch_size(),
+        ];
+        debug_assert_eq!(ans.len(), Table::num_tables());
+        ans
     }
 }
 
@@ -29,6 +40,12 @@ pub enum Table {
     Keccak = 1,
 }
 
+impl Table {
+    pub(crate) fn num_tables() -> usize {
+        Table::Keccak as usize + 1
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use anyhow::Result;

evm/src/cross_table_lookup.rs

@@ -155,7 +155,7 @@ impl<'a, F: RichField + Extendable<D>, const D: usize>
     CtlCheckVars<'a, F, F::Extension, F::Extension, D>
 {
     pub(crate) fn from_proofs<C: GenericConfig<D, F = F>>(
-        proofs: &[&StarkProofWithPublicInputs<F, C, D>],
+        proofs: &[StarkProofWithPublicInputs<F, C, D>],
         cross_table_lookups: &'a [CrossTableLookup<F>],
         ctl_challenges: &'a GrandProductChallengeSet<F>,
         num_permutation_zs: &[usize],
@@ -236,7 +236,7 @@ pub(crate) fn verify_cross_table_lookups<
     const D: usize,
 >(
     cross_table_lookups: Vec<CrossTableLookup<F>>,
-    proofs: &[&StarkProofWithPublicInputs<F, C, D>],
+    proofs: &[StarkProofWithPublicInputs<F, C, D>],
     challenges: GrandProductChallengeSet<F>,
     config: &StarkConfig,
 ) -> Result<()> {

evm/src/get_challenges.rs

@@ -1,3 +1,4 @@
+use itertools::izip;
 use plonky2::field::extension_field::Extendable;
 use plonky2::fri::proof::FriProof;
 use plonky2::hash::hash_types::RichField;
@@ -23,7 +24,7 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize> A
     ) -> AllProofChallenges<F, D> {
         let mut challenger = Challenger::<F, C::Hasher>::new();
 
-        for proof in self.proofs() {
+        for proof in &self.stark_proofs {
             challenger.observe_cap(&proof.proof.trace_cap);
         }
 
@@ -31,16 +32,15 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize> A
             get_grand_product_challenge_set(&mut challenger, config.num_challenges);
 
         AllProofChallenges {
-            cpu_challenges: self.cpu_proof.get_challenges(
-                &mut challenger,
-                &all_stark.cpu_stark,
-                config,
-            ),
-            keccak_challenges: self.keccak_proof.get_challenges(
-                &mut challenger,
-                &all_stark.keccak_stark,
-                config,
-            ),
+            stark_challenges: izip!(
+                &self.stark_proofs,
+                all_stark.nums_permutation_zs(config),
+                all_stark.permutation_batch_sizes()
+            )
+            .map(|(proof, num_perm, batch_size)| {
+                proof.get_challenges(&mut challenger, (num_perm > 0, batch_size), config)
+            })
+            .collect(),
             ctl_challenges,
         }
     }
@@ -52,10 +52,11 @@ where
     C: GenericConfig<D, F = F>,
 {
     /// Computes all Fiat-Shamir challenges used in the STARK proof.
-    pub(crate) fn get_challenges<S: Stark<F, D>>(
+    pub(crate) fn get_challenges(
         &self,
         challenger: &mut Challenger<F, C::Hasher>,
-        stark: &S,
+        // Should correspond to `(stark.use_permutation_args(), stark.permutation_batch_size())`.
+        stark_permutation_info: (bool, usize),
         config: &StarkConfig,
     ) -> StarkProofChallenges<F, D> {
         let degree_bits = self.proof.recover_degree_bits(config);
@@ -76,12 +77,8 @@ where
 
         let num_challenges = config.num_challenges;
 
-        let permutation_challenge_sets = stark.uses_permutation_args().then(|| {
-            get_n_grand_product_challenge_sets(
-                challenger,
-                num_challenges,
-                stark.permutation_batch_size(),
-            )
+        let permutation_challenge_sets = stark_permutation_info.0.then(|| {
+            get_n_grand_product_challenge_sets(challenger, num_challenges, stark_permutation_info.1)
         });
 
         challenger.observe_cap(permutation_ctl_zs_cap);

evm/src/proof.rs

@@ -19,19 +19,11 @@ use crate::permutation::GrandProductChallengeSet;
 
 #[derive(Debug, Clone)]
 pub struct AllProof<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize> {
-    pub cpu_proof: StarkProofWithPublicInputs<F, C, D>,
-    pub keccak_proof: StarkProofWithPublicInputs<F, C, D>,
-}
-
-impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize> AllProof<F, C, D> {
-    pub fn proofs(&self) -> [&StarkProofWithPublicInputs<F, C, D>; 2] {
-        [&self.cpu_proof, &self.keccak_proof]
-    }
+    pub stark_proofs: Vec<StarkProofWithPublicInputs<F, C, D>>,
 }
 
 pub(crate) struct AllProofChallenges<F: RichField + Extendable<D>, const D: usize> {
-    pub cpu_challenges: StarkProofChallenges<F, D>,
-    pub keccak_challenges: StarkProofChallenges<F, D>,
+    pub stark_challenges: Vec<StarkProofChallenges<F, D>>,
     pub ctl_challenges: GrandProductChallengeSet<F>,
 }
 

evm/src/prover.rs

@@ -48,7 +48,7 @@ where
     [(); KeccakStark::<F, D>::COLUMNS]:,
     [(); KeccakStark::<F, D>::PUBLIC_INPUTS]:,
 {
-    let num_starks = Table::Keccak as usize + 1;
+    let num_starks = Table::num_tables();
     debug_assert_eq!(num_starks, trace_poly_values.len());
     debug_assert_eq!(num_starks, public_inputs.len());
 
@@ -118,10 +118,10 @@ where
         timing,
     )?;
 
-    Ok(AllProof {
-        cpu_proof,
-        keccak_proof,
-    })
+    let stark_proofs = vec![cpu_proof, keccak_proof];
+    debug_assert_eq!(stark_proofs.len(), num_starks);
+
+    Ok(AllProof { stark_proofs })
 }
 
 /// Compute proof for a single STARK table.

evm/src/verifier.rs

@@ -32,8 +32,7 @@ where
     [(); C::Hasher::HASH_SIZE]:,
 {
     let AllProofChallenges {
-        cpu_challenges,
-        keccak_challenges,
+        stark_challenges,
         ctl_challenges,
     } = all_proof.get_challenges(&all_stark, config);
 
@@ -46,7 +45,7 @@ where
     } = all_stark;
 
     let ctl_vars_per_table = CtlCheckVars::from_proofs(
-        &all_proof.proofs(),
+        &all_proof.stark_proofs,
         &cross_table_lookups,
         &ctl_challenges,
         &nums_permutation_zs,
@@ -54,22 +53,22 @@ where
 
     verify_stark_proof_with_challenges(
         cpu_stark,
-        &all_proof.cpu_proof,
-        cpu_challenges,
+        &all_proof.stark_proofs[Table::Cpu as usize],
+        &stark_challenges[Table::Cpu as usize],
         &ctl_vars_per_table[Table::Cpu as usize],
         config,
     )?;
     verify_stark_proof_with_challenges(
         keccak_stark,
-        &all_proof.keccak_proof,
-        keccak_challenges,
+        &all_proof.stark_proofs[Table::Keccak as usize],
+        &stark_challenges[Table::Keccak as usize],
         &ctl_vars_per_table[Table::Keccak as usize],
         config,
     )?;
 
     verify_cross_table_lookups(
         cross_table_lookups,
-        &all_proof.proofs(),
+        &all_proof.stark_proofs,
         ctl_challenges,
         config,
     )
@@ -83,7 +82,7 @@ pub(crate) fn verify_stark_proof_with_challenges<
 >(
     stark: S,
     proof_with_pis: &StarkProofWithPublicInputs<F, C, D>,
-    challenges: StarkProofChallenges<F, D>,
+    challenges: &StarkProofChallenges<F, D>,
     ctl_vars: &[CtlCheckVars<F, F::Extension, F::Extension, D>],
     config: &StarkConfig,
 ) -> Result<()>
@@ -134,7 +133,7 @@ where
     let permutation_data = stark.uses_permutation_args().then(|| PermutationCheckVars {
         local_zs: permutation_ctl_zs[..num_permutation_zs].to_vec(),
         next_zs: permutation_ctl_zs_right[..num_permutation_zs].to_vec(),
-        permutation_challenge_sets: challenges.permutation_challenge_sets.unwrap(),
+        permutation_challenge_sets: challenges.permutation_challenge_sets.clone().unwrap(),
     });
     eval_vanishing_poly::<F, F::Extension, F::Extension, C, S, D, D>(
         &stark,