CommitmentScheme utils

src/core/commitment_scheme/mod.rs

@@ -16,6 +16,8 @@ use crate::commitment_scheme::merkle_decommitment::MerkleDecommitment;
 use crate::commitment_scheme::merkle_tree::MerkleTree;
 use crate::core::channel::Channel;
 
+pub mod utils;
+
 /// Holds a vector for each tree, which holds a vector for each column, which holds its respective
 /// opened values.
 pub struct OpenedValues(pub Vec<ColumnVec<Vec<BaseField>>>);

src/core/commitment_scheme/utils.rs

@@ -0,0 +1,105 @@
+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())
+    }
+    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();
+        TreeVec(zip_eq(self.0, other.0).collect())
+    }
+    pub fn as_ref(&self) -> TreeVec<&T> {
+        TreeVec(self.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())
+    }
+}
+
+impl<T> TreeVec<ColumnVec<T>> {
+    pub fn map_cols<U, F: FnMut(T) -> U>(self, mut f: F) -> TreeVec<ColumnVec<U>> {
+        TreeVec(
+            self.0
+                .into_iter()
+                .map(|column| column.into_iter().map(&mut f).collect())
+                .collect(),
+        )
+    }
+    /// Zips two [TreeVec<ColumVec<T>>] with the same structure (number of columns in each tree).
+    /// The resulting [TreeVec<ColumVec<T>>] has the same structure, with each value being a tuple
+    /// of the corresponding values from the input [TreeVec<ColumVec<T>>].
+    pub fn zip_cols<U>(
+        self,
+        other: impl Into<TreeVec<ColumnVec<U>>>,
+    ) -> TreeVec<ColumnVec<(T, U)>> {
+        let other = other.into();
+        TreeVec(
+            zip_eq(self.0, other.0)
+                .map(|(column1, column2)| zip_eq(column1, column2).collect())
+                .collect(),
+        )
+    }
+    pub fn as_cols_ref(&self) -> TreeVec<ColumnVec<&T>> {
+        TreeVec(self.iter().map(|column| column.iter().collect()).collect())
+    }
+    /// Flattens the [TreeVec<ColumVec<T>>] into a single [ColumnVec] with all the columns combined.
+    pub fn flatten(self) -> ColumnVec<T> {
+        self.0.into_iter().flatten().collect()
+    }
+}
+
+impl<'a, T> From<&'a TreeVec<ColumnVec<T>>> for TreeVec<ColumnVec<&'a T>> {
+    fn from(val: &'a TreeVec<ColumnVec<T>>) -> Self {
+        val.as_cols_ref()
+    }
+}
+
+impl<T> TreeVec<ColumnVec<Vec<T>>> {
+    /// Flattens a [TreeVec<ColumVec<T>>] of [Vec]s into a single [Vec] with all the elements
+    /// combined.
+    pub fn flatten_cols(self) -> Vec<T> {
+        self.0.into_iter().flatten().flatten().collect()
+    }
+
+    // TODO(spapini): Remove after accumulating oods quotients by size.
+    /// Flattens a [TreeVec<ColumVec<T>>] of [Vec]s into a single [Vec] with all the elements
+    /// combined, in reverse order.
+    pub fn flatten_cols_rev(self) -> Vec<T> {
+        self.0.into_iter().flatten().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>>;