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,66 @@
+#[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;
+
+    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(|| {
+                Sha256MerkleTree::new(
+                    &leaf_crh_params.clone(),
+                    &two_to_one_params.clone(),
+                    &leaves,
+                )
+                .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/crh/bowe_hopwood/mod.rs

@@ -17,7 +17,7 @@ use ark_ec::{
     twisted_edwards::Projective as TEProjective, twisted_edwards::TECurveConfig, AdditiveGroup,
     CurveGroup,
 };
-use ark_ff::{biginteger::BigInteger, fields::PrimeField};
+use ark_ff::fields::PrimeField;
 use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
 use ark_std::borrow::Borrow;
 use ark_std::cfg_chunks;
@@ -82,7 +82,7 @@ impl<P: TECurveConfig, W: pedersen::Window> CRHScheme for CRH<P, W> {
             let mut c = 0;
             let mut range = F::BigInt::from(2_u64);
             while range < upper_limit {
-                range.muln(4);
+                range <<= 4;
                 c += 1;
             }
 

src/crh/injective_map/mod.rs

@@ -1,4 +1,4 @@
-use crate::{CryptoError, Error};
+use crate::Error;
 use ark_std::rand::Rng;
 use ark_std::{fmt::Debug, hash::Hash, marker::PhantomData};
 
@@ -16,15 +16,15 @@ pub mod constraints;
 pub trait InjectiveMap<C: CurveGroup> {
     type Output: Clone + Eq + Hash + Default + Debug + CanonicalSerialize + CanonicalDeserialize;
 
-    fn injective_map(ge: &C::Affine) -> Result<Self::Output, CryptoError>;
+    fn injective_map(ge: &C::Affine) -> Result<Self::Output, Error>;
 }
 
 pub struct TECompressor;
 
 impl<P: TECurveConfig> InjectiveMap<TEProjective<P>> for TECompressor {
     type Output = <P as CurveConfig>::BaseField;
 
-    fn injective_map(ge: &TEAffine<P>) -> Result<Self::Output, CryptoError> {
+    fn injective_map(ge: &TEAffine<P>) -> Result<Self::Output, Error> {
         debug_assert!(ge.is_in_correct_subgroup_assuming_on_curve());
         Ok(ge.x)
     }

src/crh/mod.rs

@@ -21,7 +21,7 @@ pub use constraints::*;
 /// Interface to CRH. Note that in this release, while all implementations of `CRH` have fixed length,
 /// variable length CRH may also implement this trait in future.
 pub trait CRHScheme {
-    type Input: ?Sized;
+    type Input: ?Sized + Send;
     type Output: Clone
         + Eq
         + core::fmt::Debug

src/lib.rs

@@ -43,22 +43,29 @@ pub mod snark;
 #[cfg(feature = "sponge")]
 pub mod sponge;
 
-pub type Error = Box<dyn ark_std::error::Error>;
-
 #[derive(Debug)]
-pub enum CryptoError {
+pub enum Error {
     IncorrectInputLength(usize),
     NotPrimeOrder,
+    GenericError(Box<dyn ark_std::error::Error + Send>),
+    SerializationError(ark_serialize::SerializationError),
 }
 
-impl core::fmt::Display for CryptoError {
+impl core::fmt::Display for Error {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
-        let msg = match self {
-            CryptoError::IncorrectInputLength(len) => format!("input length is wrong: {}", len),
-            CryptoError::NotPrimeOrder => "element is not prime order".to_owned(),
-        };
-        write!(f, "{}", msg)
+        match self {
+            Self::IncorrectInputLength(len) => write!(f, "incorrect input length: {len}"),
+            Self::NotPrimeOrder => write!(f, "element is not prime order"),
+            Self::GenericError(e) => write!(f, "{e}"),
+            Self::SerializationError(e) => write!(f, "{e}"),
+        }
     }
 }
 
-impl ark_std::error::Error for CryptoError {}
+impl ark_std::error::Error for Error {}
+
+impl From<ark_serialize::SerializationError> for Error {
+    fn from(e: ark_serialize::SerializationError) -> Self {
+        Self::SerializationError(e)
+    }
+}

src/merkle_tree/mod.rs

@@ -14,6 +14,9 @@ mod tests;
 #[cfg(feature = "r1cs")]
 pub mod constraints;
 
+#[cfg(feature = "parallel")]
+use rayon::prelude::*;
+
 /// Convert the hash digest in different layers by converting previous layer's output to
 /// `TargetType`, which is a `Borrow` to next layer's input.
 pub trait DigestConverter<From, To: ?Sized> {
@@ -52,15 +55,16 @@ impl<T: CanonicalSerialize> DigestConverter<T, [u8]> for ByteDigestConverter<T>
 /// * `LeafHash`: Convert leaf to leaf digest
 /// * `TwoToOneHash`: Compress two inner digests to one inner digest
 pub trait Config {
-    type Leaf: ?Sized; // merkle tree does not store the leaf
-                       // leaf layer
+    type Leaf: ?Sized + Send; // merkle tree does not store the leaf
+                              // leaf layer
     type LeafDigest: Clone
         + Eq
         + core::fmt::Debug
         + 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.
@@ -229,32 +234,30 @@ impl<P: Config> MerkleTree<P> {
         height: usize,
     ) -> Result<Self, crate::Error> {
         // use empty leaf digest
-        let leaves_digest = vec![P::LeafDigest::default(); 1 << (height - 1)];
-        Self::new_with_leaf_digest(leaf_hash_param, two_to_one_hash_param, leaves_digest)
+        let leaf_digests = vec![P::LeafDigest::default(); 1 << (height - 1)];
+        Self::new_with_leaf_digest(leaf_hash_param, two_to_one_hash_param, leaf_digests)
     }
 
     /// Returns a new merkle tree. `leaves.len()` should be power of two.
-    pub fn new<L: Borrow<P::Leaf>>(
+    pub fn new<L: AsRef<P::Leaf> + Send>(
         leaf_hash_param: &LeafParam<P>,
         two_to_one_hash_param: &TwoToOneParam<P>,
-        leaves: impl IntoIterator<Item = L>,
+        #[cfg(not(feature = "parallel"))] leaves: impl IntoIterator<Item = L>,
+        #[cfg(feature = "parallel")] leaves: impl IntoParallelIterator<Item = L>,
     ) -> Result<Self, crate::Error> {
-        let mut leaves_digests = Vec::new();
-
-        // compute and store hash values for each leaf
-        for leaf in leaves.into_iter() {
-            leaves_digests.push(P::LeafHash::evaluate(leaf_hash_param, leaf)?)
-        }
+        let leaf_digests: Vec<_> = cfg_into_iter!(leaves)
+            .map(|input| P::LeafHash::evaluate(leaf_hash_param, input.as_ref()))
+            .collect::<Result<Vec<_>, _>>()?;
 
-        Self::new_with_leaf_digest(leaf_hash_param, two_to_one_hash_param, leaves_digests)
+        Self::new_with_leaf_digest(leaf_hash_param, two_to_one_hash_param, leaf_digests)
     }
 
     pub fn new_with_leaf_digest(
         leaf_hash_param: &LeafParam<P>,
         two_to_one_hash_param: &TwoToOneParam<P>,
-        leaves_digest: Vec<P::LeafDigest>,
+        leaf_digests: Vec<P::LeafDigest>,
     ) -> Result<Self, crate::Error> {
-        let leaf_nodes_size = leaves_digest.len();
+        let leaf_nodes_size = leaf_digests.len();
         assert!(
             leaf_nodes_size.is_power_of_two() && leaf_nodes_size > 1,
             "`leaves.len() should be power of two and greater than one"
@@ -266,7 +269,7 @@ impl<P: Config> MerkleTree<P> {
         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)
+        let mut non_leaf_nodes: Vec<P::InnerDigest> = cfg_into_iter!(0..non_leaf_nodes_size)
             .map(|_| hash_of_empty.clone())
             .collect();
 
@@ -282,39 +285,67 @@ impl<P: Config> MerkleTree<P> {
         {
             let start_index = level_indices.pop().unwrap();
             let upper_bound = left_child(start_index);
-            for current_index in start_index..upper_bound {
-                // `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.
-                let left_leaf_index = left_child(current_index) - upper_bound;
-                let right_leaf_index = right_child(current_index) - upper_bound;
-                // compute hash
-                non_leaf_nodes[current_index] = P::TwoToOneHash::evaluate(
-                    &two_to_one_hash_param,
-                    P::LeafInnerDigestConverter::convert(leaves_digest[left_leaf_index].clone())?,
-                    P::LeafInnerDigestConverter::convert(leaves_digest[right_leaf_index].clone())?,
-                )?
-            }
+
+            cfg_iter_mut!(non_leaf_nodes[start_index..upper_bound])
+                .enumerate()
+                .try_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;
+
+                    *n = P::TwoToOneHash::evaluate(
+                        two_to_one_hash_param,
+                        P::LeafInnerDigestConverter::convert(
+                            leaf_digests[left_leaf_index].clone(),
+                        )?,
+                        P::LeafInnerDigestConverter::convert(
+                            leaf_digests[right_leaf_index].clone(),
+                        )?,
+                    )?;
+                    Ok::<(), crate::Error>(())
+                })?;
         }
 
         // 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);
-            for current_index in start_index..upper_bound {
-                let left_index = left_child(current_index);
-                let right_index = right_child(current_index);
-                non_leaf_nodes[current_index] = P::TwoToOneHash::compress(
-                    &two_to_one_hash_param,
-                    non_leaf_nodes[left_index].clone(),
-                    non_leaf_nodes[right_index].clone(),
-                )?
-            }
-        }
 
+            let (nodes_at_level, nodes_at_prev_level) =
+                non_leaf_nodes[..].split_at_mut(upper_bound);
+            // Iterate over the nodes at the current level, and compute the hash of each node
+            cfg_iter_mut!(nodes_at_level[start_index..])
+                .enumerate()
+                .try_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;
+
+                    // need for unwrap as Box<Error> does not implement trait Send
+                    *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(),
+                    )?;
+                    Ok::<_, crate::Error>(())
+                })?;
+        }
         Ok(MerkleTree {
-            leaf_nodes: leaves_digest,
+            leaf_nodes: leaf_digests,
             non_leaf_nodes,
             height: tree_height,
             leaf_hash_param: leaf_hash_param.clone(),

src/merkle_tree/tests/constraints.rs

@@ -61,12 +61,8 @@ mod byte_mt_tests {
 
         let leaf_crh_params = <LeafH as CRHScheme>::setup(&mut rng).unwrap();
         let two_to_one_crh_params = <CompressH as TwoToOneCRHScheme>::setup(&mut rng).unwrap();
-        let mut tree = JubJubMerkleTree::new(
-            &leaf_crh_params,
-            &two_to_one_crh_params,
-            leaves.iter().map(|v| v.as_slice()),
-        )
-        .unwrap();
+        let mut tree =
+            JubJubMerkleTree::new(&leaf_crh_params, &two_to_one_crh_params, leaves).unwrap();
         let root = tree.root();
         for (i, leaf) in leaves.iter().enumerate() {
             let cs = ConstraintSystem::<Fq>::new_ref();
@@ -288,12 +284,7 @@ mod field_mt_tests {
     ) {
         let leaf_crh_params = poseidon_parameters();
         let two_to_one_params = leaf_crh_params.clone();
-        let mut tree = FieldMT::new(
-            &leaf_crh_params,
-            &two_to_one_params,
-            leaves.iter().map(|x| x.as_slice()),
-        )
-        .unwrap();
+        let mut tree = FieldMT::new(&leaf_crh_params, &two_to_one_params, leaves).unwrap();
         let root = tree.root();
         for (i, leaf) in leaves.iter().enumerate() {
             let cs = ConstraintSystem::<F>::new_ref();