From 8e220ac623ea745a30e3775c383f77e409851177 Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Tue, 23 Aug 2022 23:20:20 -0700
Subject: [PATCH] Fix for CTL challenges

See this line -

```rust
challenges: ctl_data.challenges.challenges[i % config.num_challenges],
```

This doesn't work if we have multiple lookers from the same table; then `zs_columns` will contain multiple contiguous entries for the same challenge.

We could fix the index calculation, but it seems a bit error-prone. Seems easier to store the specific challenge as part of `zs_columns`.
---
 evm/src/cross_table_lookup.rs | 55 +++++++++++++++++++----------------
 evm/src/prover.rs             | 36 ++++++++++-------------
 2 files changed, 46 insertions(+), 45 deletions(-)

diff --git a/evm/src/cross_table_lookup.rs b/evm/src/cross_table_lookup.rs
index 4097df7b..b4b8d6fb 100644
--- a/evm/src/cross_table_lookup.rs
+++ b/evm/src/cross_table_lookup.rs
@@ -168,23 +168,21 @@ impl<F: Field> CrossTableLookup<F> {
 }
 
 /// Cross-table lookup data for one table.
-#[derive(Clone)]
+#[derive(Clone, Default)]
 pub struct CtlData<F: Field> {
-    /// Challenges used in the argument.
-    pub(crate) challenges: GrandProductChallengeSet<F>,
-    /// Vector of `(Z, columns, filter_columns)` where `Z` is a Z-polynomial for a lookup
-    /// on columns `columns` with filter columns `filter_columns`.
-    pub zs_columns: Vec<(PolynomialValues<F>, Vec<Column<F>>, Option<Column<F>>)>,
+    pub(crate) zs_columns: Vec<CtlZData<F>>,
+}
+
+/// Cross-table lookup data associated with one Z(x) polynomial.
+#[derive(Clone)]
+pub(crate) struct CtlZData<F: Field> {
+    pub(crate) z: PolynomialValues<F>,
+    pub(crate) challenge: GrandProductChallenge<F>,
+    pub(crate) columns: Vec<Column<F>>,
+    pub(crate) filter_column: Option<Column<F>>,
 }
 
 impl<F: Field> CtlData<F> {
-    pub(crate) fn new(challenges: GrandProductChallengeSet<F>) -> Self {
-        Self {
-            challenges,
-            zs_columns: vec![],
-        }
-    }
-
     pub fn len(&self) -> usize {
         self.zs_columns.len()
     }
@@ -194,7 +192,10 @@ impl<F: Field> CtlData<F> {
     }
 
     pub fn z_polys(&self) -> Vec<PolynomialValues<F>> {
-        self.zs_columns.iter().map(|(p, _, _)| p.clone()).collect()
+        self.zs_columns
+            .iter()
+            .map(|zs_columns| zs_columns.z.clone())
+            .collect()
     }
 }
 
@@ -205,7 +206,7 @@ pub fn cross_table_lookup_data<F: RichField, C: GenericConfig<D, F = F>, const D
     challenger: &mut Challenger<F, C::Hasher>,
 ) -> Vec<CtlData<F>> {
     let challenges = get_grand_product_challenge_set(challenger, config.num_challenges);
-    let mut ctl_data_per_table = vec![CtlData::new(challenges.clone()); trace_poly_values.len()];
+    let mut ctl_data_per_table = vec![CtlData::default(); trace_poly_values.len()];
     for CrossTableLookup {
         looking_tables,
         looked_table,
@@ -252,19 +253,23 @@ pub fn cross_table_lookup_data<F: RichField, C: GenericConfig<D, F = F>, const D
             );
 
             for (table, z) in looking_tables.iter().zip(zs_looking) {
-                ctl_data_per_table[table.table as usize].zs_columns.push((
-                    z,
-                    table.columns.clone(),
-                    table.filter_column.clone(),
-                ));
+                ctl_data_per_table[table.table as usize]
+                    .zs_columns
+                    .push(CtlZData {
+                        z,
+                        challenge,
+                        columns: table.columns.clone(),
+                        filter_column: table.filter_column.clone(),
+                    });
             }
             ctl_data_per_table[looked_table.table as usize]
                 .zs_columns
-                .push((
-                    z_looked,
-                    looked_table.columns.clone(),
-                    looked_table.filter_column.clone(),
-                ));
+                .push(CtlZData {
+                    z: z_looked,
+                    challenge,
+                    columns: looked_table.columns.clone(),
+                    filter_column: looked_table.filter_column.clone(),
+                });
         }
     }
     ctl_data_per_table
diff --git a/evm/src/prover.rs b/evm/src/prover.rs
index 15b83173..e93ad754 100644
--- a/evm/src/prover.rs
+++ b/evm/src/prover.rs
@@ -422,17 +422,15 @@ where
                 .zs_columns
                 .iter()
                 .enumerate()
-                .map(
-                    |(i, (_, columns, filter_column))| CtlCheckVars::<F, F, P, 1> {
-                        local_z: permutation_ctl_zs_commitment.get_lde_values_packed(i_start, step)
-                            [num_permutation_zs + i],
-                        next_z: permutation_ctl_zs_commitment
-                            .get_lde_values_packed(i_next_start, step)[num_permutation_zs + i],
-                        challenges: ctl_data.challenges.challenges[i % config.num_challenges],
-                        columns,
-                        filter_column,
-                    },
-                )
+                .map(|(i, zs_columns)| CtlCheckVars::<F, F, P, 1> {
+                    local_z: permutation_ctl_zs_commitment.get_lde_values_packed(i_start, step)
+                        [num_permutation_zs + i],
+                    next_z: permutation_ctl_zs_commitment.get_lde_values_packed(i_next_start, step)
+                        [num_permutation_zs + i],
+                    challenges: zs_columns.challenge,
+                    columns: &zs_columns.columns,
+                    filter_column: &zs_columns.filter_column,
+                })
                 .collect::<Vec<_>>();
             eval_vanishing_poly::<F, F, P, C, S, D, 1>(
                 stark,
@@ -547,15 +545,13 @@ fn check_constraints<'a, F, C, S, const D: usize>(
                 .zs_columns
                 .iter()
                 .enumerate()
-                .map(
-                    |(iii, (_, columns, filter_column))| CtlCheckVars::<F, F, F, 1> {
-                        local_z: permutation_ctl_zs_subgroup_evals[i][num_permutation_zs + iii],
-                        next_z: permutation_ctl_zs_subgroup_evals[i_next][num_permutation_zs + iii],
-                        challenges: ctl_data.challenges.challenges[iii % config.num_challenges],
-                        columns,
-                        filter_column,
-                    },
-                )
+                .map(|(iii, zs_columns)| CtlCheckVars::<F, F, F, 1> {
+                    local_z: permutation_ctl_zs_subgroup_evals[i][num_permutation_zs + iii],
+                    next_z: permutation_ctl_zs_subgroup_evals[i_next][num_permutation_zs + iii],
+                    challenges: zs_columns.challenge,
+                    columns: &zs_columns.columns,
+                    filter_column: &zs_columns.filter_column,
+                })
                 .collect::<Vec<_>>();
             eval_vanishing_poly::<F, F, F, C, S, D, 1>(
                 stark,