diff --git a/crypto/src/merkle_tree/merkle.rs b/crypto/src/merkle_tree/merkle.rs index c56c8ebe6..42c4494ac 100644 --- a/crypto/src/merkle_tree/merkle.rs +++ b/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::(&mut inner_nodes, ROOT); + build::(&mut nodes, leaves_len); MerkleTree { - root: inner_nodes[ROOT].clone(), - nodes: inner_nodes, + root: nodes[ROOT].clone(), + nodes, } } diff --git a/crypto/src/merkle_tree/utils.rs b/crypto/src/merkle_tree/utils.rs index a93df7df8..87cb9b9bb 100644 --- a/crypto/src/merkle_tree/utils.rs +++ b/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(nodes: &mut Vec, 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(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::(nodes, left_child_index); - build::(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 = (1..9).map(FE::new).collect(); + let leaves_len = leaves.len(); let mut nodes = vec![FE::zero(); leaves.len() - 1]; nodes.extend(leaves); - build::>(&mut nodes, ROOT); + build::>(&mut nodes, leaves_len); assert_eq!(nodes[ROOT], FE::new(10)); } }