diff --git a/src/core/commitment_scheme/mod.rs b/src/core/commitment_scheme/mod.rs index b011bc0e0..1459e27bd 100644 --- a/src/core/commitment_scheme/mod.rs +++ b/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>>); diff --git a/src/core/commitment_scheme/utils.rs b/src/core/commitment_scheme/utils.rs new file mode 100644 index 000000000..d14598e1d --- /dev/null +++ b/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(pub Vec); +impl TreeVec { + pub fn new() -> TreeVec { + 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>(&'a self, f: F) -> TreeColumns { + TreeColumns(TreeVec(self.0.iter().map(f).collect())) + } + pub fn map U>(self, f: F) -> TreeVec { + TreeVec(self.0.into_iter().map(f).collect()) + } + pub fn zip(self, other: impl Into>) -> 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 to TreeVec<&T>. +impl<'a, T> From<&'a TreeVec> for TreeVec<&'a T> { + fn from(val: &'a TreeVec) -> Self { + val.as_ref() + } +} +impl Deref for TreeVec { + type Target = Vec; + fn deref(&self) -> &Self::Target { + &self.0 + } +} +impl DerefMut for TreeVec { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.0 + } +} +impl Default for TreeVec { + 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(pub TreeVec>); +impl TreeColumns { + pub fn new(values: Vec>) -> Self { + TreeColumns(TreeVec(values)) + } + pub fn map U>(self, mut f: F) -> TreeColumns { + 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(self, other: impl Into>) -> 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 { + self.0 .0.into_iter().flatten().collect() + } +} +impl<'a, T> From<&'a TreeColumns> for TreeColumns<&'a T> { + fn from(val: &'a TreeColumns) -> Self { + val.as_ref() + } +} +impl TreeColumns> { + /// Flattens a [TreeColumns] of [Vec]s into a single [Vec] with all the elements combined. + pub fn flatten_all(self) -> Vec { + 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 { + self.flatten().into_iter().flatten().rev().collect() + } +} diff --git a/src/core/mod.rs b/src/core/mod.rs index 4841fbd8a..6bf765838 100644 --- a/src/core/mod.rs +++ b/src/core/mod.rs @@ -28,6 +28,12 @@ impl ComponentVec> { } } +impl Default for ComponentVec { + fn default() -> Self { + Self(Vec::new()) + } +} + impl Deref for ComponentVec { type Target = Vec>;