Merge pull request #697 from mir-protocol/const_num_tables

`NUM_TABLES` constant

evm/src/all_stark.rs

@@ -41,28 +41,24 @@ impl<F: RichField + Extendable<D>, const D: usize> Default for AllStark<F, D> {
 }
 
 impl<F: RichField + Extendable<D>, const D: usize> AllStark<F, D> {
-    pub(crate) fn nums_permutation_zs(&self, config: &StarkConfig) -> Vec<usize> {
+    pub(crate) fn nums_permutation_zs(&self, config: &StarkConfig) -> [usize; NUM_TABLES] {
-        let ans = vec![
+        [
             self.cpu_stark.num_permutation_batches(config),
             self.keccak_stark.num_permutation_batches(config),
             self.keccak_memory_stark.num_permutation_batches(config),
             self.logic_stark.num_permutation_batches(config),
             self.memory_stark.num_permutation_batches(config),
-        ];
+        ]
-        debug_assert_eq!(ans.len(), Table::num_tables());
-        ans
     }
 
-    pub(crate) fn permutation_batch_sizes(&self) -> Vec<usize> {
+    pub(crate) fn permutation_batch_sizes(&self) -> [usize; NUM_TABLES] {
-        let ans = vec![
+        [
             self.cpu_stark.permutation_batch_size(),
             self.keccak_stark.permutation_batch_size(),
             self.keccak_memory_stark.permutation_batch_size(),
             self.logic_stark.permutation_batch_size(),
             self.memory_stark.permutation_batch_size(),
-        ];
+        ]
-        debug_assert_eq!(ans.len(), Table::num_tables());
-        ans
     }
 }
 
@@ -75,11 +71,7 @@ pub enum Table {
     Memory = 4,
 }
 
-impl Table {
+pub(crate) const NUM_TABLES: usize = Table::Memory as usize + 1;
-    pub(crate) fn num_tables() -> usize {
-        Table::Memory as usize + 1
-    }
-}
 
 #[allow(unused)] // TODO: Should be used soon.
 pub(crate) fn all_cross_table_lookups<F: Field>() -> Vec<CrossTableLookup<F>> {
@@ -695,7 +687,7 @@ mod tests {
             &mut memory_trace,
         );
 
-        let traces = vec![
+        let traces = [
             cpu_trace,
             keccak_trace,
             keccak_memory_trace,

evm/src/cross_table_lookup.rs

@@ -13,7 +13,7 @@ use plonky2::iop::target::Target;
 use plonky2::plonk::circuit_builder::CircuitBuilder;
 use plonky2::plonk::config::GenericConfig;
 
-use crate::all_stark::Table;
+use crate::all_stark::{Table, NUM_TABLES};
 use crate::config::StarkConfig;
 use crate::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
 use crate::permutation::{
@@ -216,12 +216,12 @@ impl<F: Field> CtlData<F> {
 
 pub fn cross_table_lookup_data<F: RichField, C: GenericConfig<D, F = F>, const D: usize>(
     config: &StarkConfig,
-    trace_poly_values: &[Vec<PolynomialValues<F>>],
+    trace_poly_values: &[Vec<PolynomialValues<F>>; NUM_TABLES],
     cross_table_lookups: &[CrossTableLookup<F>],
     challenger: &mut Challenger<F, C::Hasher>,
-) -> Vec<CtlData<F>> {
+) -> [CtlData<F>; NUM_TABLES] {
     let challenges = get_grand_product_challenge_set(challenger, config.num_challenges);
-    let mut ctl_data_per_table = vec![CtlData::default(); trace_poly_values.len()];
+    let mut ctl_data_per_table = [0; NUM_TABLES].map(|_| CtlData::default());
     for CrossTableLookup {
         looking_tables,
         looked_table,
@@ -337,12 +337,11 @@ 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: &[StarkProof<F, C, D>],
+        proofs: &[StarkProof<F, C, D>; NUM_TABLES],
         cross_table_lookups: &'a [CrossTableLookup<F>],
         ctl_challenges: &'a GrandProductChallengeSet<F>,
-        num_permutation_zs: &[usize],
+        num_permutation_zs: &[usize; NUM_TABLES],
-    ) -> Vec<Vec<Self>> {
+    ) -> [Vec<Self>; NUM_TABLES] {
-        debug_assert_eq!(proofs.len(), num_permutation_zs.len());
         let mut ctl_zs = proofs
             .iter()
             .zip(num_permutation_zs)
@@ -354,7 +353,7 @@ impl<'a, F: RichField + Extendable<D>, const D: usize>
             })
             .collect::<Vec<_>>();
 
-        let mut ctl_vars_per_table = vec![vec![]; proofs.len()];
+        let mut ctl_vars_per_table = [0; NUM_TABLES].map(|_| vec![]);
         for CrossTableLookup {
             looking_tables,
             looked_table,
@@ -441,12 +440,11 @@ pub struct CtlCheckVarsTarget<'a, F: Field, const D: usize> {
 
 impl<'a, F: Field, const D: usize> CtlCheckVarsTarget<'a, F, D> {
     pub(crate) fn from_proofs(
-        proofs: &[StarkProofTarget<D>],
+        proofs: &[StarkProofTarget<D>; NUM_TABLES],
         cross_table_lookups: &'a [CrossTableLookup<F>],
         ctl_challenges: &'a GrandProductChallengeSet<Target>,
-        num_permutation_zs: &[usize],
+        num_permutation_zs: &[usize; NUM_TABLES],
-    ) -> Vec<Vec<Self>> {
+    ) -> [Vec<Self>; NUM_TABLES] {
-        debug_assert_eq!(proofs.len(), num_permutation_zs.len());
         let mut ctl_zs = proofs
             .iter()
             .zip(num_permutation_zs)
@@ -458,7 +456,7 @@ impl<'a, F: Field, const D: usize> CtlCheckVarsTarget<'a, F, D> {
             })
             .collect::<Vec<_>>();
 
-        let mut ctl_vars_per_table = vec![vec![]; proofs.len()];
+        let mut ctl_vars_per_table = [0; NUM_TABLES].map(|_| vec![]);
         for CrossTableLookup {
             looking_tables,
             looked_table,
@@ -559,7 +557,7 @@ pub(crate) fn verify_cross_table_lookups<
     const D: usize,
 >(
     cross_table_lookups: Vec<CrossTableLookup<F>>,
-    proofs: &[StarkProof<F, C, D>],
+    proofs: &[StarkProof<F, C, D>; NUM_TABLES],
     challenges: GrandProductChallengeSet<F>,
     config: &StarkConfig,
 ) -> Result<()> {
@@ -617,7 +615,7 @@ pub(crate) fn verify_cross_table_lookups_circuit<
 >(
     builder: &mut CircuitBuilder<F, D>,
     cross_table_lookups: Vec<CrossTableLookup<F>>,
-    proofs: &[StarkProofTarget<D>],
+    proofs: &[StarkProofTarget<D>; NUM_TABLES],
     challenges: GrandProductChallengeSet<Target>,
     inner_config: &StarkConfig,
 ) {

evm/src/generation/mod.rs

@@ -4,7 +4,8 @@ use plonky2::field::polynomial::PolynomialValues;
 use plonky2::field::types::Field;
 use plonky2::hash::hash_types::RichField;
 
-use crate::all_stark::AllStark;
+use crate::all_stark::{AllStark, NUM_TABLES};
+use crate::config::StarkConfig;
 use crate::cpu::bootstrap_kernel::generate_bootstrap_kernel;
 use crate::cpu::columns::NUM_CPU_COLUMNS;
 use crate::cpu::kernel::global_metadata::GlobalMetadata;
@@ -45,7 +46,8 @@ pub struct GenerationInputs {
 pub(crate) fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
     all_stark: &AllStark<F, D>,
     inputs: GenerationInputs,
-) -> (Vec<Vec<PolynomialValues<F>>>, PublicValues) {
+    config: &StarkConfig,
+) -> ([Vec<PolynomialValues<F>>; NUM_TABLES], PublicValues) {
     let mut state = GenerationState::<F>::default();
 
     generate_bootstrap_kernel::<F>(&mut state);
@@ -83,6 +85,7 @@ pub(crate) fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
         current_cpu_row,
         memory,
         keccak_inputs,
+        keccak_memory_inputs,
         logic_ops,
         ..
     } = state;
@@ -90,9 +93,18 @@ pub(crate) fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
 
     let cpu_trace = trace_rows_to_poly_values(cpu_rows);
     let keccak_trace = all_stark.keccak_stark.generate_trace(keccak_inputs);
+    let keccak_memory_trace = all_stark
+        .keccak_memory_stark
+        .generate_trace(keccak_memory_inputs, 1 << config.fri_config.cap_height);
     let logic_trace = all_stark.logic_stark.generate_trace(logic_ops);
     let memory_trace = all_stark.memory_stark.generate_trace(memory.log);
-    let traces = vec![cpu_trace, keccak_trace, logic_trace, memory_trace];
+    let traces = [
+        cpu_trace,
+        keccak_trace,
+        keccak_memory_trace,
+        logic_trace,
+        memory_trace,
+    ];
 
     let public_values = PublicValues {
         trie_roots_before,

evm/src/get_challenges.rs

@@ -1,4 +1,3 @@
-use itertools::izip;
 use plonky2::field::extension::Extendable;
 use plonky2::fri::proof::{FriProof, FriProofTarget};
 use plonky2::hash::hash_types::RichField;
@@ -32,16 +31,18 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize> A
         let ctl_challenges =
             get_grand_product_challenge_set(&mut challenger, config.num_challenges);
 
+        let num_permutation_zs = all_stark.nums_permutation_zs(config);
+        let num_permutation_batch_sizes = all_stark.permutation_batch_sizes();
+
         AllProofChallenges {
-            stark_challenges: izip!(
+            stark_challenges: std::array::from_fn(|i| {
-                &self.stark_proofs,
+                self.stark_proofs[i].get_challenges(
-                all_stark.nums_permutation_zs(config),
+                    &mut challenger,
-                all_stark.permutation_batch_sizes()
+                    num_permutation_zs[i] > 0,
-            )
+                    num_permutation_batch_sizes[i],
-            .map(|(proof, num_perm, batch_size)| {
+                    config,
-                proof.get_challenges(&mut challenger, num_perm > 0, batch_size, config)
+                )
-            })
+            }),
-            .collect(),
             ctl_challenges,
         }
     }
@@ -66,22 +67,19 @@ impl<const D: usize> AllProofTarget<D> {
         let ctl_challenges =
             get_grand_product_challenge_set_target(builder, &mut challenger, config.num_challenges);
 
+        let num_permutation_zs = all_stark.nums_permutation_zs(config);
+        let num_permutation_batch_sizes = all_stark.permutation_batch_sizes();
+
         AllProofChallengesTarget {
-            stark_challenges: izip!(
+            stark_challenges: std::array::from_fn(|i| {
-                &self.stark_proofs,
+                self.stark_proofs[i].get_challenges::<F, C>(
-                all_stark.nums_permutation_zs(config),
-                all_stark.permutation_batch_sizes()
-            )
-            .map(|(proof, num_perm, batch_size)| {
-                proof.get_challenges::<F, C>(
                     builder,
                     &mut challenger,
-                    num_perm > 0,
+                    num_permutation_zs[i] > 0,
-                    batch_size,
+                    num_permutation_batch_sizes[i],
                     config,
                 )
-            })
+            }),
-            .collect(),
             ctl_challenges,
         }
     }

evm/src/proof.rs

@@ -13,38 +13,33 @@ use plonky2::iop::ext_target::ExtensionTarget;
 use plonky2::iop::target::Target;
 use plonky2::plonk::config::GenericConfig;
 
+use crate::all_stark::NUM_TABLES;
 use crate::config::StarkConfig;
 use crate::permutation::GrandProductChallengeSet;
 
 #[derive(Debug, Clone)]
 pub struct AllProof<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize> {
-    pub stark_proofs: Vec<StarkProof<F, C, D>>,
+    pub stark_proofs: [StarkProof<F, C, D>; NUM_TABLES],
     pub public_values: PublicValues,
 }
 
 impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize> AllProof<F, C, D> {
-    pub fn degree_bits(&self, config: &StarkConfig) -> Vec<usize> {
+    pub fn degree_bits(&self, config: &StarkConfig) -> [usize; NUM_TABLES] {
-        self.stark_proofs
+        std::array::from_fn(|i| self.stark_proofs[i].recover_degree_bits(config))
-            .iter()
-            .map(|proof| proof.recover_degree_bits(config))
-            .collect()
     }
 
-    pub fn nums_ctl_zs(&self) -> Vec<usize> {
+    pub fn nums_ctl_zs(&self) -> [usize; NUM_TABLES] {
-        self.stark_proofs
+        std::array::from_fn(|i| self.stark_proofs[i].openings.ctl_zs_last.len())
-            .iter()
-            .map(|proof| proof.openings.ctl_zs_last.len())
-            .collect()
     }
 }
 
 pub(crate) struct AllProofChallenges<F: RichField + Extendable<D>, const D: usize> {
-    pub stark_challenges: Vec<StarkProofChallenges<F, D>>,
+    pub stark_challenges: [StarkProofChallenges<F, D>; NUM_TABLES],
     pub ctl_challenges: GrandProductChallengeSet<F>,
 }
 
 pub struct AllProofTarget<const D: usize> {
-    pub stark_proofs: Vec<StarkProofTarget<D>>,
+    pub stark_proofs: [StarkProofTarget<D>; NUM_TABLES],
     pub public_values: PublicValuesTarget,
 }
 
@@ -99,7 +94,7 @@ pub struct BlockMetadataTarget {
 }
 
 pub(crate) struct AllProofChallengesTarget<const D: usize> {
-    pub stark_challenges: Vec<StarkProofChallengesTarget<D>>,
+    pub stark_challenges: [StarkProofChallengesTarget<D>; NUM_TABLES],
     pub ctl_challenges: GrandProductChallengeSet<Target>,
 }
 

evm/src/prover.rs

@@ -17,7 +17,7 @@ use plonky2::util::timing::TimingTree;
 use plonky2::util::transpose;
 use plonky2_util::{log2_ceil, log2_strict};
 
-use crate::all_stark::{AllStark, Table};
+use crate::all_stark::{AllStark, Table, NUM_TABLES};
 use crate::config::StarkConfig;
 use crate::constraint_consumer::ConstraintConsumer;
 use crate::cpu::cpu_stark::CpuStark;
@@ -53,7 +53,7 @@ where
     [(); LogicStark::<F, D>::COLUMNS]:,
     [(); MemoryStark::<F, D>::COLUMNS]:,
 {
-    let (traces, public_values) = generate_traces(all_stark, inputs);
+    let (traces, public_values) = generate_traces(all_stark, inputs, config);
     prove_with_traces(all_stark, config, traces, public_values, timing)
 }
 
@@ -61,7 +61,7 @@ where
 pub(crate) fn prove_with_traces<F, C, const D: usize>(
     all_stark: &AllStark<F, D>,
     config: &StarkConfig,
-    trace_poly_values: Vec<Vec<PolynomialValues<F>>>,
+    trace_poly_values: [Vec<PolynomialValues<F>>; NUM_TABLES],
     public_values: PublicValues,
     timing: &mut TimingTree,
 ) -> Result<AllProof<F, C, D>>
@@ -75,9 +75,6 @@ where
     [(); LogicStark::<F, D>::COLUMNS]:,
     [(); MemoryStark::<F, D>::COLUMNS]:,
 {
-    let num_starks = Table::num_tables();
-    debug_assert_eq!(num_starks, trace_poly_values.len());
-
     let rate_bits = config.fri_config.rate_bits;
     let cap_height = config.fri_config.cap_height;
 
@@ -163,14 +160,13 @@ where
         timing,
     )?;
 
-    let stark_proofs = vec![
+    let stark_proofs = [
         cpu_proof,
         keccak_proof,
         keccak_memory_proof,
         logic_proof,
         memory_proof,
     ];
-    debug_assert_eq!(stark_proofs.len(), num_starks);
 
     Ok(AllProof {
         stark_proofs,

evm/src/recursive_verifier.rs

@@ -280,7 +280,7 @@ pub fn add_virtual_all_proof<F: RichField + Extendable<D>, const D: usize>(
     degree_bits: &[usize],
     nums_ctl_zs: &[usize],
 ) -> AllProofTarget<D> {
-    let stark_proofs = vec![
+    let stark_proofs = [
         add_virtual_stark_proof(
             builder,
             all_stark.cpu_stark,
@@ -317,7 +317,6 @@ pub fn add_virtual_all_proof<F: RichField + Extendable<D>, const D: usize>(
             nums_ctl_zs[Table::Memory as usize],
         ),
     ];
-    assert_eq!(stark_proofs.len(), Table::num_tables());
 
     let public_values = add_virtual_public_values(builder);
     AllProofTarget {