CommitmentScheme utils

src/core/air/component_visitors.rs

@@ -12,7 +12,6 @@ use crate::core::circle::CirclePoint;
 use crate::core::fields::qm31::SecureField;
 use crate::core::fri::CirclePolyDegreeBound;
 use crate::core::poly::circle::{CanonicCoset, CirclePoly, SecureCirclePoly};
-use crate::core::prover::LOG_BLOWUP_FACTOR;
 use crate::core::{ColumnVec, ComponentVec};
 
 pub trait AirExt: Air<CPUBackend> {
@@ -84,16 +83,16 @@ pub trait AirExt: Air<CPUBackend> {
             fn visit<C: Component<B>>(&mut self, component: &C) {
                 let trace = self.component_traces.next().unwrap();
                 let (points, values) = component.mask_points_and_values(self.point, trace);
-                self.component_points.push(points);
-                self.component_values.push(values);
+                self.component_points.0.push(points);
+                self.component_values.0.push(values);
             }
         }
 
         let mut visitor = MaskEvaluator::<CPUBackend> {
             point,
             component_traces: component_traces.iter(),
-            component_points: ComponentVec(Vec::new()),
-            component_values: ComponentVec(Vec::new()),
+            component_points: Default::default(),
+            component_values: Default::default(),
         };
         self.visit_components(&mut visitor);
         (visitor.component_points, visitor.component_values)
@@ -116,13 +115,14 @@ pub trait AirExt: Air<CPUBackend> {
                     .map(|&log_size| CanonicCoset::new(log_size))
                     .collect_vec();
                 self.points
+                    .0
                     .push(component.mask().to_points(&domains, self.point));
             }
         }
 
         let mut visitor = MaskPointsEvaluator {
             point,
-            points: ComponentVec(Vec::new()),
+            points: Default::default(),
         };
         self.visit_components(&mut visitor);
         visitor.points
@@ -155,7 +155,7 @@ pub trait AirExt: Air<CPUBackend> {
         let mut evaluator = {
             ConstraintPointEvaluator {
                 point,
-                mask_values: mask_values.iter(),
+                mask_values: mask_values.0.iter(),
                 evaluation_accumulator: PointEvaluationAccumulator::new(
                     random_coeff,
                     self.max_constraint_log_degree_bound(),
@@ -203,31 +203,22 @@ pub trait AirExt: Air<CPUBackend> {
     }
 
     /// Returns the log degree bounds of the quotient polynomials in descending order.
-    fn commitment_domains(&self) -> Vec<CanonicCoset> {
+    fn column_log_sizes(&self) -> Vec<u32> {
         struct DomainsVisitor {
-            domains: Vec<CanonicCoset>,
+            log_sizes: Vec<u32>,
         }
 
         impl<B: Backend> ComponentVisitor<B> for DomainsVisitor {
             fn visit<C: Component<B>>(&mut self, component: &C) {
-                self.domains.extend(
-                    component
-                        .trace_log_degree_bounds()
-                        .iter()
-                        .map(|&log_size| CanonicCoset::new(log_size + LOG_BLOWUP_FACTOR)),
-                );
+                self.log_sizes.extend(component.trace_log_degree_bounds());
             }
         }
 
         let mut domains_visitor = DomainsVisitor {
-            domains: Vec::new(),
+            log_sizes: Vec::new(),
         };
         self.visit_components(&mut domains_visitor);
-        // Add the composition polynomial's domain.
-        domains_visitor.domains.push(CanonicCoset::new(
-            self.max_constraint_log_degree_bound() + LOG_BLOWUP_FACTOR,
-        ));
-        domains_visitor.domains
+        domains_visitor.log_sizes
     }
 
     fn component_traces<'a>(

src/core/commitment_scheme/utils.rs

@@ -0,0 +1,116 @@
+use std::iter::zip;
+use std::ops::{Deref, DerefMut};
+
+use itertools::zip_eq;
+
+use crate::core::ColumnVec;
+
+/// A container that holds an element for each commitment tree.
+#[derive(Debug, Clone)]
+pub struct TreeVec<T>(pub Vec<T>);
+impl<T> TreeVec<T> {
+    pub fn new() -> TreeVec<T> {
+        TreeVec(Vec::new())
+    }
+    /// Creates a [TreeColumns] instance by mapping each element of the [TreeVec] to a [ColumnVec].
+    pub fn to_cols<'a, U: 'a, F: Fn(&'a T) -> ColumnVec<U>>(&'a self, f: F) -> TreeColumns<U> {
+        TreeColumns(TreeVec(self.0.iter().map(f).collect()))
+    }
+    pub fn map<U, F: Fn(T) -> U>(self, f: F) -> TreeVec<U> {
+        TreeVec(self.0.into_iter().map(f).collect())
+    }
+    pub fn zip<U>(self, other: impl Into<TreeVec<U>>) -> TreeVec<(T, U)> {
+        let other = other.into();
+        assert_eq!(self.0.len(), other.len());
+        TreeVec(zip(self.0, other.0).collect())
+    }
+    pub fn as_ref(&self) -> TreeVec<&T> {
+        TreeVec(self.0.iter().collect())
+    }
+}
+// Converts &TreeVec<T> to TreeVec<&T>.
+impl<'a, T> From<&'a TreeVec<T>> for TreeVec<&'a T> {
+    fn from(val: &'a TreeVec<T>) -> Self {
+        val.as_ref()
+    }
+}
+impl<T> Deref for TreeVec<T> {
+    type Target = Vec<T>;
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+impl<T> DerefMut for TreeVec<T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+impl<T> Default for TreeVec<T> {
+    fn default() -> Self {
+        TreeVec(Vec::new())
+    }
+}
+
+/// A container that holds an element for each column of each commitment tree.
+#[derive(Debug, Clone)]
+pub struct TreeColumns<T>(pub TreeVec<ColumnVec<T>>);
+impl<T> TreeColumns<T> {
+    pub fn new(values: Vec<Vec<T>>) -> Self {
+        TreeColumns(TreeVec(values))
+    }
+    pub fn map<U, F: FnMut(T) -> U>(self, mut f: F) -> TreeColumns<U> {
+        TreeColumns(TreeVec(
+            self.0
+                 .0
+                .into_iter()
+                .map(|column| column.into_iter().map(&mut f).collect())
+                .collect(),
+        ))
+    }
+    /// Zips two [TreeColumns] with the same structure (number of columns in each tree).
+    /// The resulting [TreeColumns] has the same structure, with each value being a tuple of the
+    /// corresponding values from the input [TreeColumns].
+    pub fn zip_cols<U>(self, other: impl Into<TreeColumns<U>>) -> TreeColumns<(T, U)> {
+        let other = other.into();
+        assert_eq!(self.0.len(), other.0.len());
+        TreeColumns(TreeVec(
+            self.0
+                 .0
+                .into_iter()
+                .zip(other.0 .0)
+                .map(|(column1, column2)| zip_eq(column1, column2).collect())
+                .collect(),
+        ))
+    }
+    pub fn as_ref(&self) -> TreeColumns<&T> {
+        TreeColumns(TreeVec(
+            self.0
+                 .0
+                .iter()
+                .map(|column| column.iter().collect())
+                .collect(),
+        ))
+    }
+    /// Flattens the [TreeColumns] into a single [ColumnVec] with all the columns combined.
+    pub fn flatten(self) -> ColumnVec<T> {
+        self.0 .0.into_iter().flatten().collect()
+    }
+}
+impl<'a, T> From<&'a TreeColumns<T>> for TreeColumns<&'a T> {
+    fn from(val: &'a TreeColumns<T>) -> Self {
+        val.as_ref()
+    }
+}
+impl<T> TreeColumns<Vec<T>> {
+    /// Flattens a [TreeColumns] of [Vec]s into a single [Vec] with all the elements combined.
+    pub fn flatten_all(self) -> Vec<T> {
+        self.flatten().into_iter().flatten().collect()
+    }
+
+    // TODO(spapini): Remove after accumulating oods quotients by size.
+    /// Flattens a [TreeColumns] of [Vec]s into a single [Vec] with all the elements combined, in
+    /// reverse order.
+    pub fn flatten_all_rev(self) -> Vec<T> {
+        self.flatten().into_iter().flatten().rev().collect()
+    }
+}

src/core/mod.rs

@@ -28,6 +28,12 @@ impl<T: Copy> ComponentVec<ColumnVec<T>> {
     }
 }
 
+impl<T> Default for ComponentVec<T> {
+    fn default() -> Self {
+        Self(Vec::new())
+    }
+}
+
 impl<T> Deref for ComponentVec<T> {
     type Target = Vec<ColumnVec<T>>;
 

src/core/prover/mod.rs

@@ -169,7 +169,15 @@ pub fn verify(proof: StarkProof, air: &impl Air<CPUBackend>, channel: &mut Chann
         .collect_vec();
     commitment_scheme.verify(&proof.decommitments, &opening_positions);
 
-    let commitment_domains = air.commitment_domains();
+    let mut commitment_domains = air
+        .column_log_sizes()
+        .iter()
+        .map(|&log_size| CanonicCoset::new(log_size + LOG_BLOWUP_FACTOR))
+        .collect_vec();
+    commitment_domains.push(CanonicCoset::new(
+        air.max_constraint_log_degree_bound() + LOG_BLOWUP_FACTOR,
+    ));
+
     // Prepare the quotient evaluations needed for the FRI verifier.
     let sparse_circle_evaluations = prepare_fri_evaluations(
         opening_positions,

src/core/prover/utils.rs

@@ -0,0 +1,36 @@
+use crate::core::air::{Air, Component, ComponentVisitor};
+use crate::core::backend::CPUBackend;
+use crate::core::commitment_scheme::TreeColumns;
+use crate::core::ComponentVec;
+
+pub fn component_wise_to_tree_wise<T>(
+    _air: &impl Air<CPUBackend>,
+    values: ComponentVec<T>,
+) -> TreeColumns<T> {
+    TreeColumns::new(vec![values.0.into_iter().flatten().collect()])
+}
+
+pub fn tree_wise_to_component_wise<T>(
+    air: &impl Air<CPUBackend>,
+    mut values: TreeColumns<T>,
+) -> ComponentVec<T> {
+    // Recombine the trace values by component, and not by tree.
+    struct Visitor<'a, T> {
+        by_tree: &'a mut std::vec::IntoIter<T>,
+        by_component: ComponentVec<T>,
+    }
+    impl<T> ComponentVisitor<CPUBackend> for Visitor<'_, T> {
+        fn visit<C: Component<CPUBackend>>(&mut self, component: &C) {
+            self.by_component
+                .0
+                .push(self.by_tree.take(component.mask().len()).collect());
+        }
+    }
+    let values = values.0.remove(0);
+    let mut recombiner = Visitor {
+        by_tree: &mut values.into_iter(),
+        by_component: ComponentVec::default(),
+    };
+    air.visit_components(&mut recombiner);
+    recombiner.by_component
+}