Merkle trees: Parallelize building of inner levels (#824)

* implement parallel merkle level-building

* clippy

* fix typos

* restore old test case

---------

Co-authored-by: Diego K <[email protected]>

crypto/src/merkle_tree/merkle.rs

@@ -36,19 +36,19 @@ where
 
         //The leaf must be a power of 2 set
         hashed_leaves = complete_until_power_of_two(&mut hashed_leaves);
+        let leaves_len = hashed_leaves.len();
 
         //The length of leaves minus one inner node in the merkle tree
-
-        // The first elements are overwritten by build function, it doesn't matter what it's there
-        let mut inner_nodes = vec![hashed_leaves[0].clone(); hashed_leaves.len() - 1];
-        inner_nodes.extend(hashed_leaves);
+        //The first elements are overwritten by build function, it doesn't matter what it's there
+        let mut nodes = vec![hashed_leaves[0].clone(); leaves_len - 1];
+        nodes.extend(hashed_leaves);
 
         //Build the inner nodes of the tree
-        build::<B>(&mut inner_nodes, ROOT);
+        build::<B>(&mut nodes, leaves_len);
 
         MerkleTree {
-            root: inner_nodes[ROOT].clone(),
-            nodes: inner_nodes,
+            root: nodes[ROOT].clone(),
+            nodes,
         }
     }
 

crypto/src/merkle_tree/utils.rs

@@ -1,6 +1,8 @@
 use alloc::vec::Vec;
 
 use super::traits::IsMerkleTreeBackend;
+#[cfg(feature = "parallel")]
+use rayon::prelude::*;
 
 pub fn sibling_index(node_index: usize) -> usize {
     if node_index % 2 == 0 {
@@ -30,33 +32,41 @@ pub fn is_power_of_two(x: usize) -> bool {
     (x != 0) && ((x & (x - 1)) == 0)
 }
 
-pub fn build<B: IsMerkleTreeBackend>(nodes: &mut Vec<B::Node>, parent_index: usize)
+// ! CAUTION !
+// Make sure n=nodes.len()+1 is a power of two, and the last n/2 elements (leaves) are populated with hashes.
+// This function takes no precautions for other cases.
+pub fn build<B: IsMerkleTreeBackend>(nodes: &mut [B::Node], leaves_len: usize)
 where
     B::Node: Clone,
 {
-    if is_leaf(nodes.len(), parent_index) {
-        return;
+    let mut level_begin_index = leaves_len - 1;
+    let mut level_end_index = 2 * level_begin_index;
+    loop {
+        let new_level_begin_index = level_begin_index / 2;
+        let new_level_length = level_begin_index - new_level_begin_index;
+
+        let (new_level_iter, children_iter) =
+            nodes[new_level_begin_index..level_end_index + 1].split_at_mut(new_level_length);
+
+        #[cfg(feature = "parallel")]
+        let parent_and_children_zipped_iter = new_level_iter
+            .into_par_iter()
+            .zip(children_iter.par_chunks_exact(2));
+        #[cfg(not(feature = "parallel"))]
+        let parent_and_children_zipped_iter =
+            new_level_iter.iter_mut().zip(children_iter.chunks_exact(2));
+
+        parent_and_children_zipped_iter.for_each(|(new_parent, children)| {
+            *new_parent = B::hash_new_parent(&children[0], &children[1]);
+        });
+
+        level_end_index = level_begin_index - 1;
+        level_begin_index = new_level_begin_index;
+
+        if level_begin_index == level_end_index {
+            return;
+        }
     }
-
-    let left_child_index = left_child_index(parent_index);
-    let right_child_index = right_child_index(parent_index);
-
-    build::<B>(nodes, left_child_index);
-    build::<B>(nodes, right_child_index);
-
-    nodes[parent_index] = B::hash_new_parent(&nodes[left_child_index], &nodes[right_child_index]);
-}
-
-pub fn is_leaf(lenght: usize, node_index: usize) -> bool {
-    (node_index >= (lenght / 2)) && node_index < lenght
-}
-
-pub fn left_child_index(parent_index: usize) -> usize {
-    parent_index * 2 + 1
-}
-
-pub fn right_child_index(parent_index: usize) -> usize {
-    parent_index * 2 + 2
 }
 
 #[cfg(test)]
@@ -100,13 +110,14 @@ mod tests {
 
     #[test]
     // expected |10|10|13|3|7|11|2|1|2|3|4|5|6|7|8|
-    fn compleate_a_merkle_tree_from_a_set_of_leaves() {
+    fn complete_a_merkle_tree_from_a_set_of_leaves() {
         let leaves: Vec<FE> = (1..9).map(FE::new).collect();
+        let leaves_len = leaves.len();
 
         let mut nodes = vec![FE::zero(); leaves.len() - 1];
         nodes.extend(leaves);
 
-        build::<TestBackend<U64PF>>(&mut nodes, ROOT);
+        build::<TestBackend<U64PF>>(&mut nodes, leaves_len);
         assert_eq!(nodes[ROOT], FE::new(10));
     }
 }