integrated mixed degree merkle tree (#515)

src/commitment_scheme/blake2_hash.rs

@@ -64,7 +64,7 @@ impl super::hasher::Name for Blake2sHash {
 impl super::hasher::Hash<u8> for Blake2sHash {}
 
 // Wrapper for the blake2s Hashing functionalities.
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Default)]
 pub struct Blake2sHasher {
     state: Blake2s256,
 }

src/commitment_scheme/blake3_hash.rs

@@ -57,7 +57,7 @@ impl Name for Blake3Hash {
 impl super::hasher::Hash<u8> for Blake3Hash {}
 
 // Wrapper for the blake3 Hashing functionalities.
-#[derive(Clone)]
+#[derive(Clone, Default)]
 pub struct Blake3Hasher {
     state: blake3::Hasher,
 }

src/commitment_scheme/hasher.rs

@@ -22,7 +22,7 @@ pub trait Name {
 /// assert_eq!(hash, Blake3Hasher::hash(&[1, 2, 3, 4, 5, 6]));
 /// ```
 
-pub trait Hasher: Sized {
+pub trait Hasher: Sized + Default {
     type Hash: Hash<Self::NativeType>;
     type NativeType: Sized + Eq;
 

src/commitment_scheme/merkle_input.rs

@@ -117,6 +117,7 @@ impl<'a, F: Field> MerkleTreeInput<'a, F> {
 /// The column layout of a mixed degree merkle tree.
 /// The sizes of columns assigned to every layer, ordered as they were inserted & injected into hash
 /// blocks.
+#[derive(Debug, Default)]
 pub struct MerkleTreeColumnLayout {
     column_sizes: Vec<usize>,
     injected_depths_map: Vec<Vec<usize>>,
@@ -148,6 +149,13 @@ impl MerkleTreeColumnLayout {
     fn column_indices_at(&self, depth: usize) -> &[usize] {
         &self.injected_depths_map[depth - 1]
     }
+
+    pub fn build_input<'a, F: Field>(&self, columns: &[&'a [F]]) -> MerkleTreeInput<'a, F> {
+        MerkleTreeInput {
+            columns_to_inject: columns.to_vec(),
+            injected_depths_map: self.injected_depths_map.clone(),
+        }
+    }
 }
 
 #[cfg(test)]

src/commitment_scheme/mixed_degree_decommitment.rs

@@ -13,45 +13,41 @@ use crate::core::fields::{Field, IntoSlice};
 /// A correctly generated decommitment should hold all the information needed to generate the root
 /// of the tree, proving the queried values and the tree's column layout.
 // TODO(Ohad): write printing functions.
+#[derive(Debug, Default)]
 pub struct MixedDecommitment<F: Field, H: Hasher> {
     pub hashes: Vec<H::Hash>,
     pub witness_elements: Vec<F>,
 
-    // TODO(Ohad): remove this in non-debug builds.
+    // TODO(Ohad): remove these in non-debug builds.
     pub queried_values: Vec<F>,
+    pub column_layout: MerkleTreeColumnLayout,
 }
 
-#[allow(clippy::new_without_default)]
 impl<F: Field, H: Hasher> MixedDecommitment<F, H> {
     pub fn new() -> Self {
-        Self {
-            hashes: vec![],
-            witness_elements: vec![],
-            queried_values: vec![],
-        }
+        Self::default()
     }
 
     pub fn verify(
         &self,
         root: H::Hash,
-        column_layout: &MerkleTreeColumnLayout,
         queries: &[Vec<usize>],
         mut queried_values: impl Iterator<Item = F>,
     ) -> bool
     where
         F: IntoSlice<H::NativeType>,
     {
         let mut witness_hashes = self.hashes.iter();
-        let sorted_queries_by_layer = column_layout.sort_queries_by_layer(queries);
+        let sorted_queries_by_layer = self.column_layout.sort_queries_by_layer(queries);
 
         let mut next_layer_hashes = vec![];
         let mut ancestor_indices = vec![];
         let mut witness_elements = self.witness_elements.iter().copied();
-        for i in (1..=column_layout.height()).rev() {
+        for i in (1..=self.column_layout.height()).rev() {
             (next_layer_hashes, ancestor_indices) = Self::verify_single_layer(
                 i,
                 &sorted_queries_by_layer[i - 1],
-                column_layout,
+                &self.column_layout,
                 ancestor_indices.iter().copied().peekable(),
                 queried_values.by_ref(),
                 &mut witness_elements,
@@ -166,14 +162,12 @@ mod tests {
         input.insert_column(TREE_HEIGHT, &column_length_8);
         input.insert_column(TREE_HEIGHT - 1, &column_length_4);
         input.insert_column(TREE_HEIGHT - 1, &column_length_8);
-        let config = input.column_layout();
         let (tree, commitment) = MixedDegreeMerkleTree::<M31, Blake3Hasher>::commit_default(&input);
         let queries: Vec<Vec<usize>> = vec![vec![2], vec![0_usize], vec![4, 7]];
 
         let decommitment = tree.decommit(&input, &queries);
         assert!(decommitment.verify(
             commitment,
-            &config,
             &queries,
             decommitment.queried_values.iter().copied(),
         ));
@@ -189,7 +183,6 @@ mod tests {
         input.insert_column(TREE_HEIGHT, &column_length_8);
         input.insert_column(TREE_HEIGHT - 1, &column_length_4);
         input.insert_column(TREE_HEIGHT - 1, &column_length_8);
-        let config = input.column_layout();
         let (tree, _) = MixedDegreeMerkleTree::<M31, Blake3Hasher>::commit_default(&input);
         let false_commitment = Blake3Hasher::hash(b"false_commitment");
 
@@ -198,7 +191,6 @@ mod tests {
 
         assert!(decommitment.verify(
             false_commitment,
-            &config,
             &queries,
             decommitment.queried_values.iter().copied(),
         ));
@@ -214,7 +206,6 @@ mod tests {
         input.insert_column(TREE_HEIGHT, &column_length_8);
         input.insert_column(TREE_HEIGHT - 1, &column_length_4);
         input.insert_column(TREE_HEIGHT - 1, &column_length_8);
-        let config = input.column_layout();
         let (tree, commitment) = MixedDegreeMerkleTree::<M31, Blake3Hasher>::commit_default(&input);
 
         let queries: Vec<Vec<usize>> = vec![vec![2], vec![0_usize], vec![4, 7]];
@@ -223,7 +214,6 @@ mod tests {
 
         assert!(decommitment.verify(
             commitment,
-            &config,
             &queries,
             decommitment.queried_values.iter().copied(),
         ));

src/commitment_scheme/mixed_degree_merkle_tree.rs

@@ -179,6 +179,7 @@ where
                 &mut decommitment,
             );
         });
+        decommitment.column_layout = input.column_layout();
         decommitment
     }
 
@@ -641,7 +642,6 @@ mod tests {
         input.insert_column(TREE_HEIGHT - 6, &column_2);
         input.insert_column(TREE_HEIGHT - 4, &column_1);
         input.insert_column(TREE_HEIGHT, &column_3);
-        let configuration = input.column_layout();
         let (tree, commitment) = MixedDegreeMerkleTree::<M31, Blake3Hasher>::commit_default(&input);
         let queries: Vec<Vec<usize>> = vec![
             vec![2],
@@ -656,7 +656,6 @@ mod tests {
         let test_decommitment = tree.decommit(&input, &queries);
         assert!(test_decommitment.verify(
             commitment,
-            &configuration,
             &queries,
             test_decommitment.queried_values.iter().copied()
         ));

src/core/commitment_scheme/mod.rs

@@ -5,9 +5,7 @@
 //! the unique decoding regime. This is enough for a STARK proof though, where we only want to imply
 //! the existence of such polynomials, and are ok with having a small decoding list.
 //! Note: Opened points cannot come from the commitment domain.
-
 pub mod utils;
-
 use std::iter::zip;
 use std::ops::Deref;
 
@@ -34,8 +32,9 @@ use super::prover::{
 use super::queries::SparseSubCircleDomain;
 use super::ColumnVec;
 use crate::commitment_scheme::blake2_hash::{Blake2sHash, Blake2sHasher};
-use crate::commitment_scheme::merkle_decommitment::MerkleDecommitment;
-use crate::commitment_scheme::merkle_tree::MerkleTree;
+use crate::commitment_scheme::merkle_input::{MerkleTreeColumnLayout, MerkleTreeInput};
+use crate::commitment_scheme::mixed_degree_decommitment::MixedDecommitment;
+use crate::commitment_scheme::mixed_degree_merkle_tree::MixedDegreeMerkleTree;
 use crate::core::channel::Channel;
 
 type MerkleHasher = Blake2sHasher;
@@ -145,7 +144,7 @@ impl CommitmentSchemeProver {
 #[derive(Debug)]
 pub struct CommitmentSchemeProof {
     pub proved_values: TreeVec<ColumnVec<Vec<SecureField>>>,
-    pub decommitments: TreeVec<MerkleDecommitment<BaseField, MerkleHasher>>,
+    pub decommitments: TreeVec<MixedDecommitment<BaseField, MerkleHasher>>,
     pub queried_values: TreeVec<ColumnVec<Vec<BaseField>>>,
     pub proof_of_work: ProofOfWorkProof,
     pub fri_proof: FriProof<MerkleHasher>,
@@ -156,8 +155,8 @@ pub struct CommitmentSchemeProof {
 pub struct CommitmentTreeProver {
     pub polynomials: ColumnVec<CPUCirclePoly>,
     pub evaluations: ColumnVec<CPUCircleEvaluation<BaseField, BitReversedOrder>>,
-    // TODO(AlonH): Change to mixed degree merkle and remove values clone.
-    commitment: MerkleTree<BaseField, MerkleHasher>,
+    pub commitment: MixedDegreeMerkleTree<BaseField, Blake2sHasher>,
+    column_layout: MerkleTreeColumnLayout,
 }
 
 impl CommitmentTreeProver {
@@ -174,18 +173,33 @@ impl CommitmentTreeProver {
                 )
             })
             .collect_vec();
-        let commitment = MerkleTree::<BaseField, MerkleHasher>::commit(
-            evaluations
-                .iter()
-                .map(|eval| eval.values.clone())
-                .collect_vec(),
+
+        let mut merkle_input = MerkleTreeInput::new();
+        const LOG_N_BASEFIELD_ELEMENTS_IN_SACK: u32 = 4;
+
+        // The desired depth for column of log_length n is such that Blake2s hashes are filled(64B).
+        // Explicitly: There are 2^(d-1) hash 'sacks' at depth d, hence,  with elements of 4 bytes,
+        //  2^(d-1) = 2^n / 16, => d = n-3.
+        // Assuming rectangle trace, all columns go to the same depth.
+        // TOOD(AlonH): remove this assumption.
+        let inject_depth = std::cmp::max::<i32>(
+            evaluations[0].len().ilog2() as i32 - (LOG_N_BASEFIELD_ELEMENTS_IN_SACK as i32 - 1),
+            1,
         );
-        channel.mix_digest(commitment.root());
+        for column in evaluations.iter().map(|eval| &eval.values) {
+            merkle_input.insert_column(inject_depth as usize, column);
+        }
+        let (tree, root) =
+            MixedDegreeMerkleTree::<BaseField, Blake2sHasher>::commit_default(&merkle_input);
+        channel.mix_digest(root);
+
+        let column_layout = merkle_input.column_layout();
 
         CommitmentTreeProver {
             polynomials,
             evaluations,
-            commitment,
+            commitment: tree,
+            column_layout,
         }
     }
 
@@ -196,20 +210,35 @@ impl CommitmentTreeProver {
         queries: Vec<usize>,
     ) -> (
         ColumnVec<Vec<BaseField>>,
-        MerkleDecommitment<BaseField, MerkleHasher>,
+        MixedDecommitment<BaseField, Blake2sHasher>,
     ) {
         let values = self
             .evaluations
             .iter()
             .map(|c| queries.iter().map(|p| c[*p]).collect())
             .collect();
-        let decommitment = self.commitment.generate_decommitment(queries);
+        // Assuming rectangle trace, queries should be similar for all columns.
+        // TOOD(AlonH): remove this assumption.
+        let queries = std::iter::repeat(queries.to_vec())
+            .take(self.evaluations.len())
+            .collect_vec();
+
+        // Rebuild the merkle input for now.
+        // TODO(Ohad): change after tree refactor. Consider removing the input struct and have the
+        // decommitment take queries and columns only.
+        let eval_vec = self
+            .evaluations
+            .iter()
+            .map(|eval| &eval.values[..])
+            .collect_vec();
+        let input = self.column_layout.build_input(&eval_vec);
+        let decommitment = self.commitment.decommit(&input, &queries);
         (values, decommitment)
     }
 }
 
 impl Deref for CommitmentTreeProver {
-    type Target = MerkleTree<BaseField, MerkleHasher>;
+    type Target = MixedDegreeMerkleTree<BaseField, Blake2sHasher>;
 
     fn deref(&self) -> &Self::Target {
         &self.commitment
@@ -277,10 +306,18 @@ impl CommitmentSchemeVerifier {
             .as_ref()
             .zip(&proof.decommitments)
             .map(|(tree, decommitment)| {
-                // TODO(spapini): Also very proved_values here.
+                // TODO(spapini): Also verify proved_values here.
+                // Assuming columns are of equal lengths, replicate queries for all columns.
+                // TOOD(AlonH): remove this assumption.
                 tree.verify(
                     decommitment,
-                    &fri_query_domains[&(tree.log_sizes[0] + LOG_BLOWUP_FACTOR)].flatten(),
+                    &std::iter::repeat(
+                        fri_query_domains[&(tree.log_sizes[0] + LOG_BLOWUP_FACTOR)]
+                            .flatten()
+                            .clone(),
+                    )
+                    .take(tree.log_sizes.len())
+                    .collect_vec(),
                 )
             })
             .iter()
@@ -370,9 +407,13 @@ impl CommitmentTreeVerifier {
 
     pub fn verify(
         &self,
-        decommitment: &MerkleDecommitment<BaseField, MerkleHasher>,
-        positions: &[usize],
+        decommitment: &MixedDecommitment<BaseField, Blake2sHasher>,
+        queries: &[Vec<usize>],
     ) -> bool {
-        decommitment.verify(self.commitment, positions)
+        decommitment.verify(
+            self.commitment,
+            queries,
+            decommitment.queried_values.iter().copied(),
+        )
     }
 }

src/hash_functions/poseidon.rs

@@ -112,6 +112,7 @@ impl From<PoseidonHash> for [BaseField; POSEIDON_CAPACITY] {
 
 impl hasher::Hash<BaseField> for PoseidonHash {}
 
+#[derive(Default)]
 pub struct PoseidonHasher {
     params: PoseidonParams,
     state: [BaseField; POSEIDON_WIDTH],