Use simple merkle tree

src/commitment_scheme/blake2_merkle.rs

@@ -1,15 +1,15 @@
 use itertools::Itertools;
 use num_traits::Zero;
 
+use super::blake2_hash::Blake2sHash;
 use super::blake2s_ref::compress;
 use super::ops::{MerkleHasher, MerkleOps};
 use crate::core::backend::CPUBackend;
 use crate::core::fields::m31::BaseField;
 
-#[derive(Copy, Clone, PartialEq, Eq, Default)]
-pub struct Blake2sHash(pub [u32; 8]);
-pub struct Blake2Hasher;
-impl MerkleHasher for Blake2Hasher {
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Default)]
+pub struct Blake2sMerkleHasher;
+impl MerkleHasher for Blake2sMerkleHasher {
     type Hash = Blake2sHash;
 
     fn hash_node(
@@ -35,29 +35,19 @@ impl MerkleHasher for Blake2Hasher {
         for chunk in padded_values.array_chunks::<16>() {
             state = compress(state, unsafe { std::mem::transmute(chunk) }, 0, 0, 0, 0);
         }
-        Blake2sHash(state)
+        state.map(|x| x.to_le_bytes()).flatten().into()
     }
 }
 
-impl std::fmt::Debug for Blake2sHash {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        // Write as hex.
-        for &byte in self.0.iter() {
-            write!(f, "{:02x}", byte)?;
-        }
-        Ok(())
-    }
-}
-
-impl MerkleOps<Blake2Hasher> for CPUBackend {
+impl MerkleOps<Blake2sMerkleHasher> for CPUBackend {
     fn commit_on_layer(
         log_size: u32,
         prev_layer: Option<&Vec<Blake2sHash>>,
         columns: &[&Vec<BaseField>],
     ) -> Vec<Blake2sHash> {
         (0..(1 << log_size))
             .map(|i| {
-                Blake2Hasher::hash_node(
+                Blake2sMerkleHasher::hash_node(
                     prev_layer.map(|prev_layer| (prev_layer[2 * i], prev_layer[2 * i + 1])),
                     &columns.iter().map(|column| column[i]).collect_vec(),
                 )
@@ -75,17 +65,17 @@ mod tests {
     use rand::rngs::StdRng;
     use rand::{Rng, SeedableRng};
 
-    use crate::commitment_scheme::blake2_merkle::{Blake2Hasher, Blake2sHash};
+    use crate::commitment_scheme::blake2_merkle::{Blake2sHash, Blake2sMerkleHasher};
     use crate::commitment_scheme::prover::{MerkleDecommitment, MerkleProver};
     use crate::commitment_scheme::verifier::{MerkleVerificationError, MerkleVerifier};
     use crate::core::backend::CPUBackend;
     use crate::core::fields::m31::BaseField;
 
     type TestData = (
         BTreeMap<u32, Vec<usize>>,
-        MerkleDecommitment<Blake2Hasher>,
+        MerkleDecommitment<Blake2sMerkleHasher>,
         Vec<Vec<BaseField>>,
-        MerkleVerifier<Blake2Hasher>,
+        MerkleVerifier<Blake2sMerkleHasher>,
     );
     fn prepare_merkle() -> TestData {
         const N_COLS: usize = 400;
@@ -104,7 +94,8 @@ mod tests {
                     .collect_vec()
             })
             .collect_vec();
-        let merkle = MerkleProver::<CPUBackend, Blake2Hasher>::commit(cols.iter().collect_vec());
+        let merkle =
+            MerkleProver::<CPUBackend, Blake2sMerkleHasher>::commit(cols.iter().collect_vec());
 
         let mut queries = BTreeMap::<u32, Vec<usize>>::new();
         for log_size in log_size_range.rev() {
@@ -135,7 +126,7 @@ mod tests {
     #[test]
     fn test_merkle_invalid_witness() {
         let (queries, mut decommitment, values, verifier) = prepare_merkle();
-        decommitment.hash_witness[20] = Blake2sHash([0; 8]);
+        decommitment.hash_witness[20] = Blake2sHash::default();
 
         assert_eq!(
             verifier.verify(queries, values, decommitment).unwrap_err(),
@@ -190,7 +181,7 @@ mod tests {
     #[test]
     fn test_merkle_witness_too_long() {
         let (queries, mut decommitment, values, verifier) = prepare_merkle();
-        decommitment.hash_witness.push(Blake2sHash([0; 8]));
+        decommitment.hash_witness.push(Blake2sHash::default());
 
         assert_eq!(
             verifier.verify(queries, values, decommitment).unwrap_err(),

src/commitment_scheme/ops.rs

@@ -1,3 +1,5 @@
+use std::fmt::Debug;
+
 use crate::core::backend::{Col, ColumnOps};
 use crate::core::fields::m31::BaseField;
 
@@ -8,7 +10,7 @@ use crate::core::fields::m31::BaseField;
 /// children hashes.
 /// At each layer, the tree may have multiple columns of the same length as the layer.
 /// Each node in that layer contains one value from each column.
-pub trait MerkleHasher {
+pub trait MerkleHasher: Debug {
     type Hash: Clone + Eq + std::fmt::Debug;
     /// Hashes a single Merkle node. See [MerkleHasher] for more details.
     fn hash_node(

src/commitment_scheme/verifier.rs

@@ -16,6 +16,12 @@ pub struct MerkleVerifier<H: MerkleHasher> {
     pub column_log_sizes: Vec<u32>,
 }
 impl<H: MerkleHasher> MerkleVerifier<H> {
+    pub fn new(root: H::Hash, column_log_sizes: Vec<u32>) -> Self {
+        Self {
+            root,
+            column_log_sizes,
+        }
+    }
     /// Verifies the decommitment of the columns.
     ///
     /// # Arguments

src/core/commitment_scheme/mod.rs

@@ -13,4 +13,4 @@ mod verifier;
 
 pub use self::prover::{CommitmentSchemeProof, CommitmentSchemeProver};
 pub use self::utils::TreeVec;
-pub use self::verifier::{CommitmentSchemeVerifier, CommitmentTreeVerifier};
+pub use self::verifier::CommitmentSchemeVerifier;

src/core/commitment_scheme/prover.rs

@@ -1,5 +1,4 @@
-use std::iter::zip;
-use std::ops::Deref;
+use std::collections::BTreeMap;
 
 use itertools::Itertools;
 
@@ -20,13 +19,12 @@ use super::super::ColumnVec;
 use super::quotients::{compute_fri_quotients, PointSample};
 use super::utils::TreeVec;
 use crate::commitment_scheme::blake2_hash::{Blake2sHash, Blake2sHasher};
-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::commitment_scheme::blake2_merkle::Blake2sMerkleHasher;
+use crate::commitment_scheme::prover::{MerkleDecommitment, MerkleProver};
 use crate::core::channel::Channel;
 use crate::core::poly::circle::SecureEvaluation;
 
-type MerkleHasher = Blake2sHasher;
+type MerkleHasher = Blake2sMerkleHasher;
 type ProofChannel = Blake2sChannel;
 
 /// The prover side of a FRI polynomial commitment scheme. See [super].
@@ -49,7 +47,7 @@ impl CommitmentSchemeProver {
     }
 
     pub fn roots(&self) -> TreeVec<Blake2sHash> {
-        self.trees.as_ref().map(|tree| tree.root())
+        self.trees.as_ref().map(|tree| tree.commitment.root())
     }
 
     pub fn polynomials(&self) -> TreeVec<ColumnVec<&CPUCirclePoly>> {
@@ -101,7 +99,7 @@ impl CommitmentSchemeProver {
         // Run FRI commitment phase on the oods quotients.
         let fri_config = FriConfig::new(LOG_LAST_LAYER_DEGREE_BOUND, LOG_BLOWUP_FACTOR, N_QUERIES);
         let fri_prover =
-            FriProver::<CPUBackend, MerkleHasher>::commit(channel, fri_config, &quotients);
+            FriProver::<CPUBackend, Blake2sHasher>::commit(channel, fri_config, &quotients);
 
         // Proof of work.
         let proof_of_work = ProofOfWork::new(PROOF_OF_WORK_BITS).prove(channel);
@@ -111,12 +109,9 @@ impl CommitmentSchemeProver {
 
         // Decommit the FRI queries on the merkle trees.
         let decommitment_results = self.trees.as_ref().map(|tree| {
-            let queries = tree
-                .polynomials
+            let queries = fri_query_domains
                 .iter()
-                .map(|poly| {
-                    fri_query_domains[&(poly.log_size() + self.log_blowup_factor)].flatten()
-                })
+                .map(|(&log_size, domain)| (log_size, domain.flatten()))
                 .collect();
             tree.decommit(queries)
         });
@@ -137,19 +132,18 @@ impl CommitmentSchemeProver {
 #[derive(Debug)]
 pub struct CommitmentSchemeProof {
     pub sampled_values: TreeVec<ColumnVec<Vec<SecureField>>>,
-    pub decommitments: TreeVec<MixedDecommitment<BaseField, MerkleHasher>>,
+    pub decommitments: TreeVec<MerkleDecommitment<MerkleHasher>>,
     pub queried_values: TreeVec<ColumnVec<Vec<BaseField>>>,
     pub proof_of_work: ProofOfWorkProof,
-    pub fri_proof: FriProof<MerkleHasher>,
+    pub fri_proof: FriProof<Blake2sHasher>,
 }
 
 /// Prover data for a single commitment tree in a commitment scheme. The commitment scheme allows to
 /// commit on a set of polynomials at a time. This corresponds to such a set.
 pub struct CommitmentTreeProver {
     pub polynomials: ColumnVec<CPUCirclePoly>,
     pub evaluations: ColumnVec<CPUCircleEvaluation<BaseField, BitReversedOrder>>,
-    pub commitment: MixedDegreeMerkleTree<BaseField, Blake2sHasher>,
-    column_layout: MerkleTreeColumnLayout,
+    pub commitment: MerkleProver<CPUBackend, MerkleHasher>,
 }
 
 impl CommitmentTreeProver {
@@ -167,64 +161,29 @@ impl CommitmentTreeProver {
             })
             .collect_vec();
 
-        let mut merkle_input = MerkleTreeInput::new();
-        const LOG_N_BASE_FIELD_ELEMENTS_IN_SACK: u32 = 4;
-
-        for eval in evaluations.iter() {
-            // The desired depth for a 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.
-            let inject_depth = std::cmp::max::<i32>(
-                eval.domain.log_size() as i32 - (LOG_N_BASE_FIELD_ELEMENTS_IN_SACK as i32 - 1),
-                1,
-            );
-            merkle_input.insert_column(inject_depth as usize, &eval.values);
-        }
-
-        let (tree, root) =
-            MixedDegreeMerkleTree::<BaseField, Blake2sHasher>::commit_default(&merkle_input);
-        channel.mix_digest(root);
-
-        let column_layout = merkle_input.column_layout();
+        let tree = MerkleProver::commit(evaluations.iter().map(|eval| &eval.values).collect());
+        channel.mix_digest(tree.root());
 
         CommitmentTreeProver {
             polynomials,
             evaluations,
             commitment: tree,
-            column_layout,
         }
     }
 
     /// Decommits the merkle tree on the given query positions.
+    /// Returns the values at the queried positions and the decommitment.
+    /// The queries are given as a mapping from the log size of the layer size to the queried
+    /// positions on each column of that size.
     fn decommit(
         &self,
-        queries: ColumnVec<Vec<usize>>,
-    ) -> (
-        ColumnVec<Vec<BaseField>>,
-        MixedDecommitment<BaseField, Blake2sHasher>,
-    ) {
-        let values = zip(&self.evaluations, &queries)
-            .map(|(column, column_queries)| column_queries.iter().map(|q| column[*q]).collect())
-            .collect();
-
-        // 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.
+        queries: BTreeMap<u32, Vec<usize>>,
+    ) -> (ColumnVec<Vec<BaseField>>, MerkleDecommitment<MerkleHasher>) {
         let eval_vec = self
             .evaluations
             .iter()
-            .map(|eval| &eval.values[..])
+            .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 = MixedDegreeMerkleTree<BaseField, Blake2sHasher>;
-
-    fn deref(&self) -> &Self::Target {
-        &self.commitment
+        self.commitment.decommit(queries, eval_vec)
     }
 }