Implementation of Parallel Merkle Tree

.gitignore

@@ -9,4 +9,4 @@ Cargo.lock
 params
 *.swp
 *.swo
-
+.vscode

Cargo.toml

@@ -108,6 +108,12 @@ path = "benches/signature.rs"
 harness = false
 required-features = [ "signature" ]
 
+[[bench]]
+name = "merkle_tree"
+path = "benches/merkle_tree.rs"
+harness = false
+required-features = [ "merkle_tree" ]
+
 [patch.crates-io]
 ark-r1cs-std = { git = "https://github.com/arkworks-rs/r1cs-std/" }
 ark-ff = { git = "https://github.com/arkworks-rs/algebra/" }

benches/merkle_tree.rs

@@ -0,0 +1,83 @@
+#[macro_use]
+extern crate criterion;
+
+static NUM_LEAVES: i32 = 1 << 20;
+
+mod bytes_mt_benches {
+    use ark_crypto_primitives::crh::*;
+    use ark_crypto_primitives::merkle_tree::*;
+    use ark_crypto_primitives::to_uncompressed_bytes;
+    use ark_ff::BigInteger256;
+    use ark_serialize::CanonicalSerialize;
+    use ark_std::{test_rng, UniformRand};
+    use criterion::Criterion;
+    use std::borrow::Borrow;
+
+    #[cfg(feature = "parallel")]
+    use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
+
+    use crate::NUM_LEAVES;
+
+    type LeafH = sha2::Sha256;
+    type CompressH = sha2::Sha256;
+
+    struct Sha256MerkleTreeParams;
+
+    impl Config for Sha256MerkleTreeParams {
+        type Leaf = [u8];
+
+        type LeafDigest = <LeafH as CRHScheme>::Output;
+        type LeafInnerDigestConverter = ByteDigestConverter<Self::LeafDigest>;
+        type InnerDigest = <CompressH as TwoToOneCRHScheme>::Output;
+
+        type LeafHash = LeafH;
+        type TwoToOneHash = CompressH;
+    }
+    type Sha256MerkleTree = MerkleTree<Sha256MerkleTreeParams>;
+
+    pub fn merkle_tree_create(c: &mut Criterion) {
+        let mut rng = test_rng();
+        let leaves: Vec<_> = (0..NUM_LEAVES)
+            .map(|_| {
+                let rnd = BigInteger256::rand(&mut rng);
+                to_uncompressed_bytes!(rnd).unwrap()
+            })
+            .collect();
+        let leaf_crh_params = <LeafH as CRHScheme>::setup(&mut rng).unwrap();
+        let two_to_one_params = <CompressH as TwoToOneCRHScheme>::setup(&mut rng)
+            .unwrap()
+            .clone();
+        c.bench_function("Merkle Tree Create (Leaves as [u8])", move |b| {
+            b.iter(|| {
+                #[cfg(not(feature = "parallel"))]
+                {
+                    _ = Sha256MerkleTree::new(
+                        &leaf_crh_params.clone(),
+                        &two_to_one_params.clone(),
+                        leaves.iter().map(|x| x.as_slice()),
+                    )
+                    .unwrap();
+                }
+                #[cfg(feature = "parallel")]
+                {
+                    _ = Sha256MerkleTree::new(
+                        &leaf_crh_params.clone(),
+                        &two_to_one_params.clone(),
+                        leaves.par_iter().map(|x| x.as_slice()),
+                    )
+                    .unwrap();
+                }
+            })
+        });
+    }
+
+    criterion_group! {
+        name = mt_create;
+        config = Criterion::default().sample_size(10);
+        targets = merkle_tree_create
+    }
+}
+
+criterion_main!(
+    crate::bytes_mt_benches::mt_create,
+);

src/merkle_tree/mod.rs

@@ -8,6 +8,9 @@ use ark_std::borrow::Borrow;
 use ark_std::hash::Hash;
 use ark_std::vec::Vec;
 
+#[cfg(feature = "parallel")]
+use rayon::{iter::IntoParallelIterator, prelude::*};
+
 #[cfg(test)]
 mod tests;
 
@@ -60,7 +63,8 @@ pub trait Config {
         + Hash
         + Default
         + CanonicalSerialize
-        + CanonicalDeserialize;
+        + CanonicalDeserialize
+        + Send;
     // transition between leaf layer to inner layer
     type LeafInnerDigestConverter: DigestConverter<
         Self::LeafDigest,
@@ -73,7 +77,8 @@ pub trait Config {
         + Hash
         + Default
         + CanonicalSerialize
-        + CanonicalDeserialize;
+        + CanonicalDeserialize
+        + Send;
 
     // Tom's Note: in the future, if we want different hash function, we can simply add more
     // types of digest here and specify a digest converter. Same for constraints.
@@ -234,26 +239,145 @@ impl<P: Config> MerkleTree<P> {
     }
 
     /// Returns a new merkle tree. `leaves.len()` should be power of two.
+    #[cfg(not(feature = "parallel"))]
     pub fn new<L: Borrow<P::Leaf>>(
         leaf_hash_param: &LeafParam<P>,
         two_to_one_hash_param: &TwoToOneParam<P>,
         leaves: impl IntoIterator<Item = L>,
     ) -> Result<Self, crate::Error> {
-        let mut leaves_digests = Vec::new();
+        // Serial version of new
+        // Compute and store hash values for each leaf
+        let leaves_digest: Result<Vec<_>, _> = leaves
+            .into_iter()
+            .map(|leaf| P::LeafHash::evaluate(leaf_hash_param, leaf))
+            .collect();
+
+        Self::new_with_leaf_digest(leaf_hash_param, two_to_one_hash_param, leaves_digest?)
+    }
+
+    /// Returns a new merkle tree. `leaves.len()` should be power of two.
+    #[cfg(feature = "parallel")]
+    pub fn new<L: Borrow<P::Leaf>>(
+        leaf_hash_param: &LeafParam<P>,
+        two_to_one_hash_param: &TwoToOneParam<P>,
+        leaves: impl IntoParallelIterator<Item = L>,
+    ) -> Result<Self, crate::Error> {
+        // Parallel version of new
+        // TODO: Box<Error> is not Send, I have used .unwrap() but this changes semantics between
+        // parallel and serial version. We should fix this in the future.
+        let leaves_digest: Vec<_> = leaves
+            .into_par_iter()
+            .map(|leaf| P::LeafHash::evaluate(leaf_hash_param, leaf).unwrap())
+            .collect();
+
+        Self::new_with_leaf_digest(leaf_hash_param, two_to_one_hash_param, leaves_digest)
+    }
+
+    #[cfg(feature = "parallel")]
+    pub fn new_with_leaf_digest(
+        leaf_hash_param: &LeafParam<P>,
+        two_to_one_hash_param: &TwoToOneParam<P>,
+        leaves_digest: Vec<P::LeafDigest>,
+    ) -> Result<Self, crate::Error> {
+        // Parallel implementation of new_with_leaf_digest
+
+        let leaf_nodes_size = leaves_digest.len();
+        assert!(
+            leaf_nodes_size.is_power_of_two() && leaf_nodes_size > 1,
+            "`leaves.len() should be power of two and greater than one"
+        );
+        let non_leaf_nodes_size = leaf_nodes_size - 1;
+
+        let tree_height = tree_height(leaf_nodes_size);
+
+        let hash_of_empty: P::InnerDigest = P::InnerDigest::default();
+
+        // initialize the merkle tree as array of nodes in level order
+        let mut non_leaf_nodes: Vec<P::InnerDigest> = (0..non_leaf_nodes_size)
+            .into_par_iter()
+            .map(|_| hash_of_empty.clone())
+            .collect();
+
+        // Compute the starting indices for each non-leaf level of the tree
+        let mut index = 0;
+        let mut level_indices = Vec::with_capacity(tree_height - 1);
+        for _ in 0..(tree_height - 1) {
+            level_indices.push(index);
+            index = left_child(index);
+        }
 
-        // compute and store hash values for each leaf
-        for leaf in leaves.into_iter() {
-            leaves_digests.push(P::LeafHash::evaluate(leaf_hash_param, leaf)?)
+        // compute the hash values for the non-leaf bottom layer
+        {
+            let start_index = level_indices.pop().unwrap();
+            let upper_bound = left_child(start_index);
+            let bottom_nodes = &mut non_leaf_nodes[start_index..upper_bound];
+
+            bottom_nodes.par_iter_mut().enumerate().for_each(|(i, n)| {
+                // `left_child(current_index)` and `right_child(current_index) returns the position of
+                // leaf in the whole tree (represented as a list in level order). We need to shift it
+                // by `-upper_bound` to get the index in `leaf_nodes` list.
+
+                //similarly, we need to rescale i by start_index
+                //to get the index outside the slice and in the level-ordered list of nodes
+
+                let current_index = i + start_index;
+                let left_leaf_index = left_child(current_index) - upper_bound;
+                let right_leaf_index = right_child(current_index) - upper_bound;
+                // compute hash
+                *n = P::TwoToOneHash::evaluate(
+                    two_to_one_hash_param,
+                    P::LeafInnerDigestConverter::convert(leaves_digest[left_leaf_index].clone())
+                        .unwrap(),
+                    P::LeafInnerDigestConverter::convert(leaves_digest[right_leaf_index].clone())
+                        .unwrap(),
+                )
+                .unwrap()
+            });
+        }
+
+        // compute the hash values for nodes in every other layer in the tree
+        level_indices.reverse();
+        for &start_index in &level_indices {
+            // The layer beginning `start_index` ends at `upper_bound` (exclusive).
+
+            let upper_bound = left_child(start_index);
+
+            let (nodes_at_level, nodes_at_prev_level) =
+                non_leaf_nodes[..].split_at_mut(upper_bound);
+            nodes_at_level[start_index..]
+                .par_iter_mut()
+                .enumerate()
+                .for_each(|(i, n)| {
+                    let current_index = i + start_index;
+                    let left_leaf_index = left_child(current_index) - upper_bound;
+                    let right_leaf_index = right_child(current_index) - upper_bound;
+
+                    *n = P::TwoToOneHash::compress(
+                        two_to_one_hash_param,
+                        nodes_at_prev_level[left_leaf_index].clone(),
+                        nodes_at_prev_level[right_leaf_index].clone(),
+                    )
+                    .unwrap()
+                });
         }
 
-        Self::new_with_leaf_digest(leaf_hash_param, two_to_one_hash_param, leaves_digests)
+        Ok(MerkleTree {
+            leaf_nodes: leaves_digest,
+            non_leaf_nodes,
+            height: tree_height,
+            leaf_hash_param: leaf_hash_param.clone(),
+            two_to_one_hash_param: two_to_one_hash_param.clone(),
+        })
     }
 
+    #[cfg(not(feature = "parallel"))]
     pub fn new_with_leaf_digest(
         leaf_hash_param: &LeafParam<P>,
         two_to_one_hash_param: &TwoToOneParam<P>,
         leaves_digest: Vec<P::LeafDigest>,
     ) -> Result<Self, crate::Error> {
+        // Serial implementation of new_with_leaf_digest
+
         let leaf_nodes_size = leaves_digest.len();
         assert!(
             leaf_nodes_size.is_power_of_two() && leaf_nodes_size > 1,