Change signature to enable correct iteration

arkworks-rs · Dec 28, 2023 · 93e0e8d · 93e0e8d
1 parent c17337f
commit 93e0e8d
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 42 deletions.
diff --git a/benches/merkle_tree.rs b/benches/merkle_tree.rs
@@ -49,7 +49,7 @@ mod bytes_mt_benches {
                 Sha256MerkleTree::new(
                     &leaf_crh_params.clone(),
                     &two_to_one_params.clone(),
-                    leaves.iter().map(|x| x.as_slice()),
+                    &leaves,
                 )
                 .unwrap();
             })

diff --git a/src/crh/mod.rs b/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

diff --git a/src/merkle_tree/mod.rs b/src/merkle_tree/mod.rs
@@ -239,21 +239,14 @@ impl<P: Config> MerkleTree<P> {
     }
 
     /// Returns a new merkle tree. `leaves.len()` should be power of two.
-    pub fn new<L, T>(
+    pub fn new<L: AsRef<P::Leaf> + Send>(
         leaf_hash_param: &LeafParam<P>,
         two_to_one_hash_param: &TwoToOneParam<P>,
-        leaves: T,
-    ) -> Result<Self, crate::Error>
-    where
-        L: Borrow<P::Leaf> + Send,
-        T: IntoIterator<Item = L>,
-        T::IntoIter: Send,
-    {
-        #[cfg(feature = "parallel")]
-        let leaves = leaves.into_iter().par_bridge();
-
+        #[cfg(not(feature = "parallel"))] leaves: impl IntoIterator<Item = L>,
+        #[cfg(feature = "parallel")] leaves: impl IntoParallelIterator<Item = L>,
+    ) -> Result<Self, crate::Error> {
         let leaf_digests: Vec<_> = cfg_into_iter!(leaves)
-            .map(|leaf| P::LeafHash::evaluate(leaf_hash_param, leaf))
+            .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, &leaf_digests)

diff --git a/src/merkle_tree/tests/constraints.rs b/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();

diff --git a/src/merkle_tree/tests/mod.rs b/src/merkle_tree/tests/mod.rs
@@ -46,16 +46,10 @@ mod bytes_mt_tests {
             .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();
-
-        let mut tree = JubJubMerkleTree::new(
-            &leaf_crh_params.clone(),
-            &two_to_one_params.clone(),
-            leaves.iter().map(|x| x.as_slice()),
-        )
-        .unwrap();
+        let two_to_one_params = <CompressH as TwoToOneCRHScheme>::setup(&mut rng).unwrap();
+
+        let mut tree =
+            JubJubMerkleTree::new(&leaf_crh_params, &two_to_one_params, &leaves).unwrap();
 
         let mut root = tree.root();
         // test merkle tree functionality without update
@@ -149,12 +143,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 mut root = tree.root();