Dumb down FRI

src/commitment_scheme/prover.rs

@@ -7,11 +7,32 @@ use crate::core::backend::{Col, Column};
 use crate::core::fields::m31::BaseField;
 
 pub struct MerkleProver<B: MerkleOps<H>, H: MerkleHasher> {
+    /// Layers of the Merkle tree.
+    /// The first layer is the largest column.
+    /// The last layer is the root.
+    /// Layer n holds the 2^n hashes of the previous layer if exists, and all the columns of size
+    /// 2^n.
     pub layers: Vec<Col<B, H::Hash>>,
 }
+/// The MerkleProver struct represents a prover for a Merkle commitment scheme.
+/// It is generic over the types `B` and `H`, which represent the Merkle operations and Merkle
+/// hasher respectively.
 impl<B: MerkleOps<H>, H: MerkleHasher> MerkleProver<B, H> {
     /// Commits to columns.
     /// Columns must be of power of 2 sizes and sorted in descending order.
+    ///
+    /// # Arguments
+    ///
+    /// * `columns` - A vector of references to columns.
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if the columns are not sorted in descending order or if the columns
+    /// vector is empty.
+    ///
+    /// # Returns
+    ///
+    /// A new instance of `MerkleProver` with the committed layers.
     pub fn commit(columns: Vec<&Col<B, BaseField>>) -> Self {
         // Check that columns are of descending order.
         assert!(!columns.is_empty());
@@ -34,9 +55,17 @@ impl<B: MerkleOps<H>, H: MerkleHasher> MerkleProver<B, H> {
 
     /// Decommits to columns on the given queries.
     /// Queries are given as indices to the largest column.
+    ///
+    /// # Arguments
+    ///
+    /// * `queries` - A vector of query indices to the largest column.
+    ///
+    /// # Returns
+    ///
+    /// A `Decommitment` struct containing the witness.
     pub fn decommit(&self, mut queries: Vec<usize>) -> Decommitment<H> {
         let mut witness = Vec::new();
-        for layer in &self.layers {
+        for layer in &self.layers[..self.layers.len() - 1] {
             let mut queries_iter = queries.into_iter().peekable();
 
             // Propagate queries and hashes to the next layer.
@@ -46,9 +75,7 @@ impl<B: MerkleOps<H>, H: MerkleHasher> MerkleProver<B, H> {
                 if queries_iter.next_if_eq(&(query ^ 1)).is_some() {
                     continue;
                 }
-                if layer.len() > 1 {
-                    witness.push(layer.at(query ^ 1));
-                }
+                witness.push(layer.at(query ^ 1));
             }
             queries = next_queries;
         }

src/commitment_scheme/verifier.rs

@@ -15,10 +15,34 @@ pub struct MerkleTreeVerifier<H: MerkleHasher> {
 }
 impl<H: MerkleHasher> MerkleTreeVerifier<H> {
     /// Verifies the decommitment of the columns.
-    /// Queries are given as indices to the largest column.
-    /// Values are given as pair of log_size of the column, and the decommited values of the
-    /// column.
-    /// Must be given in the same order as the columns were committed.
+    ///
+    /// # Arguments
+    ///
+    /// * `queries` - A vector of indices representing the queries to the largest column. This is
+    ///   assumed to be sufficient for our use case, though it could be extended to support queries
+    ///   to any column.
+    /// * `values` - A vector of pairs containing the log_size of the column and the decommitted
+    ///   values of the column. Must be given in the same order as the columns were committed.
+    /// * `decommitment` - The decommitment object containing the witness and column values.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if any of the following conditions are met:
+    ///
+    /// * The witness is too long (not fully consumed).
+    /// * The witness is too short (missing values).
+    /// * The column values are too long (not fully consumed).
+    /// * The column values are too short (missing values).
+    /// * The computed root does not match the expected root.
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if the `values` vector is not sorted in descending order based on
+    /// the `log_size` of the columns.
+    ///
+    /// # Returns
+    ///
+    /// Returns `Ok(())` if the decommitment is successfully verified.
     pub fn verify(
         &self,
         queries: Vec<usize>,
@@ -59,6 +83,19 @@ struct MerkleVerifier<H: MerkleHasher> {
     layer_column_values: Vec<std::vec::IntoIter<BaseField>>,
 }
 impl<H: MerkleHasher> MerkleVerifier<H> {
+    /// Computes the root hash of a Merkle tree from the decommitment information.
+    ///
+    /// # Arguments
+    ///
+    /// * `queries` - A vector of query indices to the largest column.
+    ///
+    /// # Returns
+    ///
+    /// Returns the computed root hash of the Merkle tree.
+    ///
+    /// # Errors
+    ///
+    /// Returns a `MerkleVerificationError` if there is an error during the computation.
     pub fn compute_root_from_decommitment(
         &mut self,
         queries: Vec<usize>,

src/core/air/accumulation.rs

@@ -135,7 +135,7 @@ impl DomainEvaluationAccumulator<CPUBackend> {
             .zip(self.n_cols_per_size.iter())
             .skip(1)
         {
-            let coeffs = SecureColumn {
+            let coeffs = SecureColumn::<CPUBackend> {
                 columns: values.columns.map(|c| {
                     CPUCircleEvaluation::<_, BitReversedOrder>::new(
                         CanonicCoset::new(log_size as u32).circle_domain(),

src/core/backend/cpu/fri.rs

@@ -1,31 +1,25 @@
-use std::iter::zip;
-
 use super::CPUBackend;
 use crate::core::fft::ibutterfly;
-use crate::core::fields::m31::BaseField;
 use crate::core::fields::qm31::SecureField;
-use crate::core::fields::{ExtensionOf, Field, FieldExpOps};
+use crate::core::fields::FieldExpOps;
 use crate::core::fri::{FriOps, CIRCLE_TO_LINE_FOLD_STEP, FOLD_STEP};
-use crate::core::poly::circle::CircleEvaluation;
+use crate::core::poly::circle::SecureEvaluation;
 use crate::core::poly::line::LineEvaluation;
-use crate::core::poly::BitReversedOrder;
 use crate::core::utils::bit_reverse_index;
 
 impl FriOps for CPUBackend {
-    fn fold_line(
-        eval: &LineEvaluation<Self, SecureField, BitReversedOrder>,
-        alpha: SecureField,
-    ) -> LineEvaluation<Self, SecureField, BitReversedOrder> {
+    fn fold_line(eval: &LineEvaluation<Self>, alpha: SecureField) -> LineEvaluation<Self> {
         let n = eval.len();
         assert!(n >= 2, "Evaluation too small");
 
         let domain = eval.domain();
 
         let folded_values = eval
             .values
+            .into_iter()
             .array_chunks()
             .enumerate()
-            .map(|(i, &[f_x, f_neg_x])| {
+            .map(|(i, [f_x, f_neg_x])| {
                 // TODO(andrew): Inefficient. Update when domain twiddles get stored in a buffer.
                 let x = domain.at(bit_reverse_index(i << FOLD_STEP, domain.log_size()));
 
@@ -37,22 +31,20 @@ impl FriOps for CPUBackend {
 
         LineEvaluation::new(domain.double(), folded_values)
     }
-    fn fold_circle_into_line<F: Field>(
-        dst: &mut LineEvaluation<Self, SecureField, BitReversedOrder>,
-        src: &CircleEvaluation<Self, F, BitReversedOrder>,
+    fn fold_circle_into_line(
+        dst: &mut LineEvaluation<Self>,
+        src: &SecureEvaluation<Self>,
         alpha: SecureField,
-    ) where
-        F: ExtensionOf<BaseField>,
-        SecureField: ExtensionOf<F> + Field,
-    {
+    ) {
         assert_eq!(src.len() >> CIRCLE_TO_LINE_FOLD_STEP, dst.len());
 
         let domain = src.domain;
         let alpha_sq = alpha * alpha;
 
-        zip(&mut dst.values, src.array_chunks())
+        src.into_iter()
+            .array_chunks()
             .enumerate()
-            .for_each(|(i, (dst, &[f_p, f_neg_p]))| {
+            .for_each(|(i, [f_p, f_neg_p])| {
                 // TODO(andrew): Inefficient. Update when domain twiddles get stored in a buffer.
                 let p = domain.at(bit_reverse_index(
                     i << CIRCLE_TO_LINE_FOLD_STEP,
@@ -64,7 +56,7 @@ impl FriOps for CPUBackend {
                 ibutterfly(&mut f0_px, &mut f1_px, p.y.inverse());
                 let f_prime = alpha * f1_px + f0_px;
 
-                *dst = *dst * alpha_sq + f_prime;
+                dst.values.set(i, dst.values.at(i) * alpha_sq + f_prime);
             });
     }
 }

src/core/backend/cpu/mod.rs

@@ -8,7 +8,6 @@ use std::fmt::Debug;
 use super::{Backend, Column, ColumnOps, FieldOps};
 use crate::core::fields::Field;
 use crate::core::poly::circle::{CircleEvaluation, CirclePoly};
-use crate::core::poly::line::LineEvaluation;
 use crate::core::utils::bit_reverse;
 
 #[derive(Copy, Clone, Debug)]
@@ -50,7 +49,6 @@ impl<T: Debug + Clone + Default> Column<T> for Vec<T> {
 pub type CPUCirclePoly = CirclePoly<CPUBackend>;
 pub type CPUCircleEvaluation<F, EvalOrder> = CircleEvaluation<CPUBackend, F, EvalOrder>;
 // TODO(spapini): Remove the EvalOrder on LineEvaluation.
-pub type CPULineEvaluation<F, EvalOrder> = LineEvaluation<CPUBackend, F, EvalOrder>;
 
 #[cfg(test)]
 mod tests {

src/core/commitment_scheme/prover.rs

@@ -98,7 +98,11 @@ impl CommitmentSchemeProver {
         let fri_prover = FriProver::<CPUBackend, MerkleHasher>::commit(
             channel,
             fri_config,
-            &quotients.flatten_cols_rev(),
+            &quotients
+                .flatten_cols_rev()
+                .into_iter()
+                .map(|e| e.into())
+                .collect_vec(),
         );
 
         // Proof of work.

src/core/commitment_scheme/verifier.rs

@@ -147,7 +147,7 @@ fn eval_quotients_on_sparse_domain(
     commitment_domain: CircleDomain,
     point: CirclePoint<SecureField>,
     value: SecureField,
-) -> Result<SparseCircleEvaluation<SecureField>, VerificationError> {
+) -> Result<SparseCircleEvaluation, VerificationError> {
     let queried_values = &mut queried_values.into_iter();
     let res = SparseCircleEvaluation::new(
         query_domains

src/core/fields/secure_column.rs

@@ -9,14 +9,11 @@ pub const SECURE_EXTENSION_DEGREE: usize =
 
 /// An array of `SECURE_EXTENSION_DEGREE` base field columns, that represents a column of secure
 /// field elements.
+#[derive(Clone, Debug)]
 pub struct SecureColumn<B: Backend> {
     pub columns: [Col<B, BaseField>; SECURE_EXTENSION_DEGREE],
 }
 impl SecureColumn<CPUBackend> {
-    pub fn at(&self, index: usize) -> SecureField {
-        SecureField::from_m31_array(std::array::from_fn(|i| self.columns[i][index]))
-    }
-
     pub fn set(&mut self, index: usize, value: SecureField) {
         self.columns
             .iter_mut()
@@ -25,6 +22,10 @@ impl SecureColumn<CPUBackend> {
     }
 }
 impl<B: Backend> SecureColumn<B> {
+    pub fn at(&self, index: usize) -> SecureField {
+        SecureField::from_m31_array(std::array::from_fn(|i| self.columns[i].at(index)))
+    }
+
     pub fn zeros(len: usize) -> Self {
         Self {
             columns: std::array::from_fn(|_| Col::<B, BaseField>::zeros(len)),
@@ -38,4 +39,56 @@ impl<B: Backend> SecureColumn<B> {
     pub fn is_empty(&self) -> bool {
         self.columns[0].is_empty()
     }
+
+    pub fn to_cpu(&self) -> SecureColumn<CPUBackend> {
+        SecureColumn {
+            columns: self.columns.clone().map(|c| c.to_vec()),
+        }
+    }
+}
+
+pub struct SecureColumnIter<'a> {
+    column: &'a SecureColumn<CPUBackend>,
+    index: usize,
+}
+impl Iterator for SecureColumnIter<'_> {
+    type Item = SecureField;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.index < self.column.len() {
+            let value = self.column.at(self.index);
+            self.index += 1;
+            Some(value)
+        } else {
+            None
+        }
+    }
+}
+impl<'a> IntoIterator for &'a SecureColumn<CPUBackend> {
+    type Item = SecureField;
+    type IntoIter = SecureColumnIter<'a>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        SecureColumnIter {
+            column: self,
+            index: 0,
+        }
+    }
+}
+impl FromIterator<SecureField> for SecureColumn<CPUBackend> {
+    fn from_iter<I: IntoIterator<Item = SecureField>>(iter: I) -> Self {
+        let mut columns = std::array::from_fn(|_| vec![]);
+        for value in iter.into_iter() {
+            let vals = value.to_m31_array();
+            for j in 0..SECURE_EXTENSION_DEGREE {
+                columns[j].push(vals[j]);
+            }
+        }
+        SecureColumn { columns }
+    }
+}
+impl From<SecureColumn<CPUBackend>> for Vec<SecureField> {
+    fn from(column: SecureColumn<CPUBackend>) -> Self {
+        column.into_iter().collect()
+    }
 }