avoid unnecessary allocation in initial Merkle tree

Co-Authored-By: Borna <94551425+Barnadrot@users.noreply.github.com>

Author: Faten848

crates/whir/src/commit.rs

@@ -9,8 +9,8 @@ use crate::*;
 
 #[derive(Debug, Clone)]
 pub enum MerkleData<EF: ExtensionField<PF<EF>>> {
-    Base(RoundMerkleTree<PF<EF>, PF<EF>>),
-    Extension(RoundMerkleTree<PF<EF>, EF>),
+    Base(RoundMerkleTree<PF<EF>>),
+    Extension(RoundMerkleTree<PF<EF>>),
 }
 
 impl<EF: ExtensionField<PF<EF>>> MerkleData<EF> {

crates/whir/src/matrix.rs

@@ -2,12 +2,11 @@
 
 use std::{
     borrow::{Borrow, BorrowMut},
-    iter,
     marker::PhantomData,
     ops::Deref,
 };
 
-use field::{ExtensionField, Field, PackedValue};
+use field::PackedValue;
 use itertools::Itertools;
 
 pub trait Matrix<T: Send + Sync + Clone>: Send + Sync {
@@ -123,88 +122,6 @@ impl<T: Clone + Send + Sync, S: DenseStorage<T>> DenseMatrix<T, S> {
     }
 }
 
-#[derive(Debug, Clone)]
-pub struct FlatMatrixView<F, EF, Inner>(Inner, PhantomData<(F, EF)>);
-
-impl<F, EF, Inner> FlatMatrixView<F, EF, Inner> {
-    pub const fn new(inner: Inner) -> Self {
-        Self(inner, PhantomData)
-    }
-}
-
-impl<F, EF, Inner> Deref for FlatMatrixView<F, EF, Inner> {
-    type Target = Inner;
-
-    fn deref(&self) -> &Self::Target {
-        &self.0
-    }
-}
-
-impl<F, EF, Inner> Matrix<F> for FlatMatrixView<F, EF, Inner>
-where
-    F: Field,
-    EF: ExtensionField<F>,
-    Inner: Matrix<EF>,
-{
-    fn width(&self) -> usize {
-        self.0.width() * EF::DIMENSION
-    }
-
-    fn height(&self) -> usize {
-        self.0.height()
-    }
-
-    unsafe fn row_subseq_unchecked(
-        &self,
-        r: usize,
-        start: usize,
-        end: usize,
-    ) -> impl IntoIterator<Item = F, IntoIter = impl Iterator<Item = F> + Send + Sync> {
-        // We can skip the first start / EF::DIMENSION elements in the row.
-        let len = end - start;
-        let inner_start = start / EF::DIMENSION;
-        unsafe {
-            // Safety: The caller must ensure that r < self.height(), start <= end and end < self.width().
-            FlatIter {
-                inner: self
-                    .0
-                    // We set end to be the width of the inner matrix and use take to ensure we get the right
-                    // number of elements.
-                    .row_subseq_unchecked(r, inner_start, self.0.width())
-                    .into_iter()
-                    .peekable(),
-                idx: start,
-                _phantom: PhantomData,
-            }
-            .take(len)
-        }
-    }
-}
-
-pub struct FlatIter<F, I: Iterator> {
-    inner: iter::Peekable<I>,
-    idx: usize,
-    _phantom: PhantomData<F>,
-}
-
-impl<F, EF, I> Iterator for FlatIter<F, I>
-where
-    F: Field,
-    EF: ExtensionField<F>,
-    I: Iterator<Item = EF>,
-{
-    type Item = F;
-    fn next(&mut self) -> Option<Self::Item> {
-        if self.idx == EF::DIMENSION {
-            self.idx = 0;
-            self.inner.next();
-        }
-        let value = self.inner.peek()?.as_basis_coefficients_slice()[self.idx];
-        self.idx += 1;
-        Some(value)
-    }
-}
-
 #[derive(Copy, Clone, PartialEq, Eq, Debug)]
 pub struct Dimensions {
     /// Number of columns in the matrix.

crates/whir/src/merkle.rs

@@ -19,58 +19,88 @@ use utils::log2_ceil_usize;
 
 use crate::DenseMatrix;
 use crate::Dimensions;
-use crate::FlatMatrixView;
 use crate::Matrix;
 pub use symetric::DIGEST_ELEMS;
 
-pub(crate) type RoundMerkleTree<F, EF> = WhirMerkleTree<F, FlatMatrixView<F, EF, DenseMatrix<EF>>, DIGEST_ELEMS>;
+pub(crate) type RoundMerkleTree<F> = WhirMerkleTree<F, DenseMatrix<F>, DIGEST_ELEMS>;
 
 #[allow(clippy::missing_transmute_annotations)]
 pub(crate) fn merkle_commit<F: Field, EF: ExtensionField<F>>(
     matrix: DenseMatrix<EF>,
     full_n_cols: usize,
     effective_n_cols: usize,
-) -> ([F; DIGEST_ELEMS], RoundMerkleTree<F, EF>) {
-    let perm = default_koalabear_poseidon1_16();
+) -> ([F; DIGEST_ELEMS], RoundMerkleTree<F>) {
     if TypeId::of::<(F, EF)>() == TypeId::of::<(KoalaBear, QuinticExtensionFieldKB)>() {
         let matrix = unsafe { std::mem::transmute::<_, DenseMatrix<QuinticExtensionFieldKB>>(matrix) };
-        let view = FlatMatrixView::new(matrix);
         let dim = <QuinticExtensionFieldKB as BasedVectorSpace<KoalaBear>>::DIMENSION;
+        let dft_base_width = matrix.width * dim;
         let full_base_width = full_n_cols * dim;
         let effective_base_width = effective_n_cols * dim;
-        let tree =
-            WhirMerkleTree::new::<PFPacking<KoalaBear>, _, 16, 8>(&perm, view, full_base_width, effective_base_width);
+        let base_values = QuinticExtensionFieldKB::flatten_to_base(matrix.values);
+        let base_matrix = DenseMatrix::<KoalaBear>::new(base_values, dft_base_width);
+        let tree = build_merkle_tree_koalabear(base_matrix, full_base_width, effective_base_width);
         let root: [_; DIGEST_ELEMS] = tree.root();
         let root = unsafe { std::mem::transmute_copy::<_, [F; DIGEST_ELEMS]>(&root) };
-        let tree = unsafe { std::mem::transmute::<_, RoundMerkleTree<F, EF>>(tree) };
+        let tree = unsafe { std::mem::transmute::<_, RoundMerkleTree<F>>(tree) };
         (root, tree)
     } else if TypeId::of::<(F, EF)>() == TypeId::of::<(KoalaBear, KoalaBear)>() {
         let matrix = unsafe { std::mem::transmute::<_, DenseMatrix<KoalaBear>>(matrix) };
-        let tree = WhirMerkleTree::new::<PFPacking<KoalaBear>, _, 16, 8>(&perm, matrix, full_n_cols, effective_n_cols);
+        let tree = build_merkle_tree_koalabear(matrix, full_n_cols, effective_n_cols);
         let root: [_; DIGEST_ELEMS] = tree.root();
         let root = unsafe { std::mem::transmute_copy::<_, [F; DIGEST_ELEMS]>(&root) };
-        let tree = unsafe { std::mem::transmute::<_, RoundMerkleTree<F, EF>>(tree) };
+        let tree = unsafe { std::mem::transmute::<_, RoundMerkleTree<F>>(tree) };
         (root, tree)
     } else {
         unimplemented!()
     }
 }
 
+#[instrument(name = "build merkle tree", skip_all)]
+fn build_merkle_tree_koalabear(
+    leaf: DenseMatrix<KoalaBear>,
+    full_base_width: usize,
+    effective_base_width: usize,
+) -> RoundMerkleTree<KoalaBear> {
+    let perm = default_koalabear_poseidon1_16();
+    let n_zero_suffix_rate_chunks = (full_base_width - effective_base_width) / 8;
+    let first_layer = if n_zero_suffix_rate_chunks >= 2 {
+        let scalar_state = symetric::precompute_zero_suffix_state::<KoalaBear, _, 16, 8, DIGEST_ELEMS>(
+            &perm,
+            n_zero_suffix_rate_chunks,
+        );
+        let packed_state: [PFPacking<KoalaBear>; 16] =
+            std::array::from_fn(|i| PFPacking::<KoalaBear>::from_fn(|_| scalar_state[i]));
+        first_digest_layer_with_initial_state::<PFPacking<KoalaBear>, _, _, DIGEST_ELEMS, 16, 8>(
+            &perm,
+            &leaf,
+            &packed_state,
+            effective_base_width,
+        )
+    } else {
+        first_digest_layer::<PFPacking<KoalaBear>, _, _, DIGEST_ELEMS, 16, 8>(&perm, &leaf, full_base_width)
+    };
+    let tree = symetric::merkle::MerkleTree::from_first_layer::<PFPacking<KoalaBear>, _, 16>(&perm, first_layer);
+    WhirMerkleTree {
+        leaf,
+        tree,
+        full_leaf_base_width: full_base_width,
+    }
+}
+
 #[allow(clippy::missing_transmute_annotations)]
 pub(crate) fn merkle_open<F: Field, EF: ExtensionField<F>>(
-    merkle_tree: &RoundMerkleTree<F, EF>,
+    merkle_tree: &RoundMerkleTree<F>,
     index: usize,
 ) -> (Vec<EF>, Vec<[F; DIGEST_ELEMS]>) {
     if TypeId::of::<(F, EF)>() == TypeId::of::<(KoalaBear, QuinticExtensionFieldKB)>() {
-        let merkle_tree =
-            unsafe { std::mem::transmute::<_, &RoundMerkleTree<KoalaBear, QuinticExtensionFieldKB>>(merkle_tree) };
+        let merkle_tree = unsafe { std::mem::transmute::<_, &RoundMerkleTree<KoalaBear>>(merkle_tree) };
         let (inner_leaf, proof) = merkle_tree.open(index);
         let leaf = QuinticExtensionFieldKB::reconstitute_from_base(inner_leaf);
         let leaf = unsafe { std::mem::transmute::<_, Vec<EF>>(leaf) };
         let proof = unsafe { std::mem::transmute::<_, Vec<[F; DIGEST_ELEMS]>>(proof) };
         (leaf, proof)
     } else if TypeId::of::<(F, EF)>() == TypeId::of::<(KoalaBear, KoalaBear)>() {
-        let merkle_tree = unsafe { std::mem::transmute::<_, &RoundMerkleTree<KoalaBear, KoalaBear>>(merkle_tree) };
+        let merkle_tree = unsafe { std::mem::transmute::<_, &RoundMerkleTree<KoalaBear>>(merkle_tree) };
         let (inner_leaf, proof) = merkle_tree.open(index);
         let leaf = KoalaBear::reconstitute_from_base(inner_leaf);
         let leaf = unsafe { std::mem::transmute::<_, Vec<EF>>(leaf) };