[wip]Sai/support optional table

starky/src/cross_table_lookup.rs

@@ -226,7 +226,7 @@ impl<'a, F: Field> CtlData<'a, F> {
 /// the CTL data necessary to prove a multi-STARK system.
 pub fn get_ctl_data<'a, F, C, const D: usize, const N: usize>(
     config: &StarkConfig,
-    trace_poly_values: &[Vec<PolynomialValues<F>>; N],
+    trace_poly_values: &[Option<Vec<PolynomialValues<F>>>; N],
     all_cross_table_lookups: &'a [CrossTableLookup<F>],
     challenger: &mut Challenger<F, C::Hasher>,
     max_constraint_degree: usize,
@@ -321,7 +321,7 @@ pub(crate) fn get_ctl_auxiliary_polys<F: Field>(
 ///
 /// For each `CrossTableLookup`, and each looking/looked table, the partial products for the CTL are computed, and added to the said table's `CtlZData`.
 pub(crate) fn cross_table_lookup_data<'a, F: RichField, const D: usize, const N: usize>(
-    trace_poly_values: &[Vec<PolynomialValues<F>>; N],
+    trace_poly_values: &[Option<Vec<PolynomialValues<F>>>; N],
     cross_table_lookups: &'a [CrossTableLookup<F>],
     ctl_challenges: &GrandProductChallengeSet<F>,
     constraint_degree: usize,
@@ -341,8 +341,15 @@ pub(crate) fn cross_table_lookup_data<'a, F: RichField, const D: usize, const N:
                 constraint_degree,
             );
 
+            let trace_values = trace_poly_values[looked_table.table]
+                .as_ref()
+                .expect(&format!(
+                    "Missing trace polynomial values for table {:?}",
+                    looked_table.table
+                ));
+
             let z_looked = partial_sums(
-                &trace_poly_values[looked_table.table],
+                trace_values,
                 &[(&looked_table.columns, &looked_table.filter)],
                 challenge,
                 constraint_degree,
@@ -394,7 +401,7 @@ pub(crate) fn cross_table_lookup_data<'a, F: RichField, const D: usize, const N:
 /// Computes helper columns and Z polynomials for all looking tables
 /// of one cross-table lookup (i.e. for one looked table).
 fn ctl_helper_zs_cols<F: Field, const N: usize>(
-    all_stark_traces: &[Vec<PolynomialValues<F>>; N],
+    all_stark_traces: &[Option<Vec<PolynomialValues<F>>>; N],
     looking_tables: Vec<TableWithColumns<F>>,
     challenge: GrandProductChallenge<F>,
     constraint_degree: usize,
@@ -404,17 +411,17 @@ fn ctl_helper_zs_cols<F: Field, const N: usize>(
     grouped_lookups
         .into_iter()
         .map(|(table, group)| {
+            let trace_values = all_stark_traces[table].as_ref().expect(&format!(
+                "Missing trace polynomial values for table {:?}",
+                table
+            ));
+
             let columns_filters = group
                 .map(|table| (&table.columns[..], &table.filter))
                 .collect::<Vec<(&[Column<F>], &Filter<F>)>>();
             (
                 table,
-                partial_sums(
-                    &all_stark_traces[table],
-                    &columns_filters,
-                    challenge,
-                    constraint_degree,
-                ),
+                partial_sums(trace_values, &columns_filters, challenge, constraint_degree),
             )
         })
         .collect::<Vec<(usize, Vec<PolynomialValues<F>>)>>()
@@ -494,7 +501,7 @@ impl<'a, F: RichField + Extendable<D>, const D: usize>
 {
     /// Extracts the `CtlCheckVars` for each STARK.
     pub fn from_proofs<C: GenericConfig<D, F = F>, const N: usize>(
-        proofs: &[StarkProofWithMetadata<F, C, D>; N],
+        proofs: &[Option<StarkProofWithMetadata<F, C, D>>; N],
         cross_table_lookups: &'a [CrossTableLookup<F>],
         ctl_challenges: &'a GrandProductChallengeSet<F>,
         num_lookup_columns: &[usize; N],
@@ -504,10 +511,14 @@ impl<'a, F: RichField + Extendable<D>, const D: usize>
         // If there are no auxiliary polys in the proofs `openings`,
         // return early. The verifier will reject the proofs when
         // calling `validate_proof_shape`.
-        if proofs
-            .iter()
-            .any(|p| p.proof.openings.auxiliary_polys.is_none())
-        {
+        if proofs.iter().any(|p| {
+            // Check only if `p` is `Some`, otherwise ignore
+            if let Some(proof) = p {
+                proof.proof.openings.auxiliary_polys.is_none()
+            } else {
+                false // If `p` is `None`, we don't want to count it, so return false
+            }
+        }) {
             return ctl_vars_per_table;
         }
 
@@ -523,21 +534,27 @@ impl<'a, F: RichField + Extendable<D>, const D: usize>
             .iter()
             .zip(num_lookup_columns)
             .map(|(p, &num_lookup)| {
-                let openings = &p.proof.openings;
-
-                let ctl_zs = &openings
-                    .auxiliary_polys
-                    .as_ref()
-                    .expect("We cannot have CTls without auxiliary polynomials.")[num_lookup..];
-                let ctl_zs_next = &openings
-                    .auxiliary_polys_next
-                    .as_ref()
-                    .expect("We cannot have CTls without auxiliary polynomials.")[num_lookup..];
-                ctl_zs.iter().zip(ctl_zs_next).collect::<Vec<_>>()
+                if let Some(proof) = p {
+                    let openings = &proof.proof.openings;
+
+                    let ctl_zs = &openings
+                        .auxiliary_polys
+                        .as_ref()
+                        .expect("We cannot have CTLs without auxiliary polynomials.")[num_lookup..];
+                    let ctl_zs_next = &openings
+                        .auxiliary_polys_next
+                        .as_ref()
+                        .expect("We cannot have CTLs without auxiliary polynomials.")[num_lookup..];
+                    Some(ctl_zs.iter().zip(ctl_zs_next).collect::<Vec<_>>())
+                } else {
+                    None // If `p` is None, return None
+                }
             })
             .collect::<Vec<_>>();
 
-        // Put each cross-table lookup polynomial into the correct table data: if a CTL polynomial is extracted from looking/looked table t, then we add it to the `CtlCheckVars` of table t.
+        // Put each cross-table lookup polynomial into the correct table data: if a CTL polynomial
+        // is extracted from looking/looked table t, then we add it to the `CtlCheckVars` of table
+        // t.
         let mut start_indices = [0; N];
         let mut z_indices = [0; N];
         for (
@@ -559,9 +576,11 @@ impl<'a, F: RichField + Extendable<D>, const D: usize>
                 }
 
                 for &table in filtered_looking_tables.iter() {
+                    let ctl_z = ctl_zs[table].as_ref().unwrap();
+
                     // We have first all the helper polynomials, then all the z polynomials.
                     let (looking_z, looking_z_next) =
-                        ctl_zs[table][total_num_helper_cols_by_table[table] + z_indices[table]];
+                        ctl_z[total_num_helper_cols_by_table[table] + z_indices[table]];
 
                     let count = looking_tables
                         .iter()
@@ -580,7 +599,7 @@ impl<'a, F: RichField + Extendable<D>, const D: usize>
                         columns.push(&col[..]);
                         filter.push(filt.clone());
                     }
-                    let helper_columns = ctl_zs[table]
+                    let helper_columns = ctl_z
                         [start_indices[table]..start_indices[table] + num_ctls[table]]
                         .iter()
                         .map(|&(h, _)| *h)
@@ -599,9 +618,10 @@ impl<'a, F: RichField + Extendable<D>, const D: usize>
                     });
                 }
 
-                let (looked_z, looked_z_next) = ctl_zs[looked_table.table]
-                    [total_num_helper_cols_by_table[looked_table.table]
-                        + z_indices[looked_table.table]];
+                let ctl_z = ctl_zs[looked_table.table].as_ref().unwrap();
+                let (looked_z, looked_z_next) = ctl_z[total_num_helper_cols_by_table
+                    [looked_table.table]
+                    + z_indices[looked_table.table]];
 
                 z_indices[looked_table.table] += 1;
 
@@ -924,11 +944,13 @@ pub(crate) fn eval_cross_table_lookup_checks_circuit<
 /// The key of `ctl_extra_looking_sums` is the corresponding CTL's position within `cross_table_lookups`.
 pub fn verify_cross_table_lookups<F: RichField + Extendable<D>, const D: usize, const N: usize>(
     cross_table_lookups: &[CrossTableLookup<F>],
-    ctl_zs_first: [Vec<F>; N],
+    ctl_zs_first: [Option<Vec<F>>; N],
     ctl_extra_looking_sums: &HashMap<usize, Vec<F>>,
     config: &StarkConfig,
 ) -> Result<()> {
-    let mut ctl_zs_openings = ctl_zs_first.iter().map(|v| v.iter()).collect::<Vec<_>>();
+    let mut ctl_zs_openings: [_; N] =
+        core::array::from_fn(|i| ctl_zs_first[i].as_ref().map(|vec| vec.iter()));
+
     for (
         index,
         CrossTableLookup {
@@ -947,16 +969,26 @@ pub fn verify_cross_table_lookups<F: RichField + Extendable<D>, const D: usize,
         }
         for c in 0..config.num_challenges {
             // Compute the combination of all looking table CTL polynomial openings.
-
             let looking_zs_sum = filtered_looking_tables
                 .iter()
-                .map(|&table| *ctl_zs_openings[table].next().unwrap())
+                .map(|&table| {
+                    *ctl_zs_openings[table]
+                        .as_mut()
+                        .expect("CTL opening should be present")
+                        .next()
+                        .expect("CTL value should be available")
+                })
                 .sum::<F>()
                 // Get elements looking into `looked_table` that are not associated to any STARK.
                 + ctl_extra_looking_sum.map(|v| v[c]).unwrap_or_default();
 
             // Get the looked table CTL polynomial opening.
-            let looked_z = *ctl_zs_openings[looked_table.table].next().unwrap();
+            let looked_z = *ctl_zs_openings[looked_table.table]
+                .as_mut()
+                .expect("CTL opening for looked table should be present")
+                .next()
+                .expect("CTL value should be available for looked table");
+
             // Ensure that the combination of looking table openings is equal to the looked table opening.
             ensure!(
                 looking_zs_sum == looked_z,
@@ -965,7 +997,10 @@ pub fn verify_cross_table_lookups<F: RichField + Extendable<D>, const D: usize,
             );
         }
     }
-    debug_assert!(ctl_zs_openings.iter_mut().all(|iter| iter.next().is_none()));
+
+    debug_assert!(ctl_zs_openings
+        .iter_mut()
+        .all(|iter| iter.as_mut().map_or(true, |it| it.next().is_none())));
 
     Ok(())
 }
@@ -979,11 +1014,13 @@ pub fn verify_cross_table_lookups_circuit<
 >(
     builder: &mut CircuitBuilder<F, D>,
     cross_table_lookups: Vec<CrossTableLookup<F>>,
-    ctl_zs_first: [Vec<Target>; N],
+    ctl_zs_first: [Option<Vec<Target>>; N],
     ctl_extra_looking_sums: &HashMap<usize, Vec<Target>>,
     inner_config: &StarkConfig,
 ) {
-    let mut ctl_zs_openings = ctl_zs_first.iter().map(|v| v.iter()).collect::<Vec<_>>();
+    let mut ctl_zs_openings: [_; N] =
+        core::array::from_fn(|i| ctl_zs_first[i].as_ref().map(|vec| vec.iter()));
+
     for (
         index,
         CrossTableLookup {
@@ -1002,23 +1039,34 @@ pub fn verify_cross_table_lookups_circuit<
         }
         for c in 0..inner_config.num_challenges {
             // Compute the combination of all looking table CTL polynomial openings.
-            let mut looking_zs_sum = builder.add_many(
-                filtered_looking_tables
-                    .iter()
-                    .map(|&table| *ctl_zs_openings[table].next().unwrap()),
-            );
+            let mut looking_zs_sum =
+                builder.add_many(filtered_looking_tables.iter().map(|&table| {
+                    *ctl_zs_openings[table]
+                        .as_mut()
+                        .expect("CTL opening should be present")
+                        .next()
+                        .expect("CTL value should be available")
+                }));
 
             // Get elements looking into `looked_table` that are not associated to any STARK.
             let extra_sum = ctl_extra_looking_sum.map(|v| v[c]).unwrap_or_default();
             looking_zs_sum = builder.add(looking_zs_sum, extra_sum);
 
             // Get the looked table CTL polynomial opening.
-            let looked_z = *ctl_zs_openings[looked_table.table].next().unwrap();
+            let looked_z = *ctl_zs_openings[looked_table.table]
+                .as_mut()
+                .expect("CTL opening for looked table should be present")
+                .next()
+                .expect("CTL value should be available for looked table");
+
             // Verify that the combination of looking table openings is equal to the looked table opening.
             builder.connect(looked_z, looking_zs_sum);
         }
     }
-    debug_assert!(ctl_zs_openings.iter_mut().all(|iter| iter.next().is_none()));
+
+    debug_assert!(ctl_zs_openings
+        .iter_mut()
+        .all(|iter| iter.as_mut().map_or(true, |it| it.next().is_none())));
 }
 
 /// Debugging module used to assert correctness of the different CTLs of a multi-STARK system,
@@ -1040,7 +1088,7 @@ pub mod debug_utils {
     /// Check that the provided traces and cross-table lookups are consistent.
     /// The key of `extra_looking_values` is the corresponding CTL's position within `cross_table_lookups`.
     pub fn check_ctls<F: Field>(
-        trace_poly_values: &[Vec<PolynomialValues<F>>],
+        trace_poly_values: &[Option<Vec<PolynomialValues<F>>>],
         cross_table_lookups: &[CrossTableLookup<F>],
         extra_looking_values: &HashMap<usize, Vec<Vec<F>>>,
     ) {
@@ -1050,7 +1098,7 @@ pub mod debug_utils {
     }
 
     fn check_ctl<F: Field>(
-        trace_poly_values: &[Vec<PolynomialValues<F>>],
+        trace_poly_values: &[Option<Vec<PolynomialValues<F>>>],
         ctl: &CrossTableLookup<F>,
         ctl_index: usize,
         extra_looking_values: Option<&Vec<Vec<F>>>,
@@ -1096,11 +1144,15 @@ pub mod debug_utils {
     }
 
     fn process_table<F: Field>(
-        trace_poly_values: &[Vec<PolynomialValues<F>>],
+        trace_poly_values: &[Option<Vec<PolynomialValues<F>>>],
         table: &TableWithColumns<F>,
         multiset: &mut MultiSet<F>,
     ) {
-        let trace = &trace_poly_values[table.table];
+        let trace = trace_poly_values[table.table].as_ref().expect(&format!(
+            "Missing trace polynomial values for table {:?}",
+            table.table
+        ));
+
         for i in 0..trace[0].len() {
             let filter = table.filter.eval_table(trace, i);
             if filter.is_one() {

starky/src/proof.rs

@@ -166,7 +166,7 @@ pub struct MultiProof<
     const N: usize,
 > {
     /// Proofs for all the different STARK modules.
-    pub stark_proofs: [StarkProofWithMetadata<F, C, D>; N],
+    pub stark_proofs: [Option<StarkProofWithMetadata<F, C, D>>; N],
     /// Cross-table lookup challenges.
     pub ctl_challenges: GrandProductChallengeSet<F>,
 }
@@ -176,8 +176,12 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize, c
 {
     /// Returns the degree (i.e. the trace length) of each STARK proof,
     /// from their common [`StarkConfig`].
-    pub fn recover_degree_bits(&self, config: &StarkConfig) -> [usize; N] {
-        core::array::from_fn(|i| self.stark_proofs[i].proof.recover_degree_bits(config))
+    pub fn recover_degree_bits(&self, config: &StarkConfig) -> [Option<usize>; N] {
+        core::array::from_fn(|i| {
+            self.stark_proofs[i]
+                .as_ref()
+                .map(|proof| proof.proof.recover_degree_bits(config))
+        })
     }
 }