From ee14c9243d43f4590890f03c78e32b25ef99a41f Mon Sep 17 00:00:00 2001 From: ljedrz Date: Mon, 25 Mar 2024 11:44:39 +0100 Subject: [PATCH 1/4] perf: reduce zero-padding in MerkleTree Signed-off-by: ljedrz --- .../src/merkle_tree/helpers/path_hash.rs | 15 ++++- console/collections/src/merkle_tree/mod.rs | 64 +++++++++++++++++-- .../src/merkle_tree/tests/append.rs | 14 ++-- 3 files changed, 73 insertions(+), 20 deletions(-) diff --git a/console/collections/src/merkle_tree/helpers/path_hash.rs b/console/collections/src/merkle_tree/helpers/path_hash.rs index fc9841672c..83167d0385 100644 --- a/console/collections/src/merkle_tree/helpers/path_hash.rs +++ b/console/collections/src/merkle_tree/helpers/path_hash.rs @@ -31,11 +31,20 @@ pub trait PathHash: Clone + Send + Sync { fn hash_children(&self, left: &Self::Hash, right: &Self::Hash) -> Result; /// Returns the hash for each tuple of child nodes. - fn hash_all_children(&self, child_nodes: &[(Self::Hash, Self::Hash)]) -> Result> { + /// If there are no children, the supplied empty hash is used. + fn hash_all_children( + &self, + child_nodes: &[Option<(Self::Hash, Self::Hash)>], + empty_node_hash: Self::Hash, + ) -> Result> { + let hash_children = |children: &Option<(Self::Hash, Self::Hash)>| { + if let Some((left, right)) = children { self.hash_children(left, right) } else { Ok(empty_node_hash) } + }; + match child_nodes.len() { 0 => Ok(vec![]), - 1..=100 => child_nodes.iter().map(|(left, right)| self.hash_children(left, right)).collect(), - _ => cfg_iter!(child_nodes).map(|(left, right)| self.hash_children(left, right)).collect(), + 1..=100 => child_nodes.iter().map(hash_children).collect(), + _ => cfg_iter!(child_nodes).map(hash_children).collect(), } } } diff --git a/console/collections/src/merkle_tree/mod.rs b/console/collections/src/merkle_tree/mod.rs index 2e3c273e63..a36e3c65c4 100644 --- a/console/collections/src/merkle_tree/mod.rs +++ b/console/collections/src/merkle_tree/mod.rs @@ -77,8 +77,12 @@ impl, PH: PathHash // Compute the empty hash. let empty_hash = path_hasher.hash_empty()?; + // Calculate the size of the tree which excludes leafless nodes. + let minimum_tree_size = + std::cmp::max(1, num_nodes + leaves.len() + if leaves.len() > 1 { leaves.len() % 2 } else { 0 }); + // Initialize the Merkle tree. - let mut tree = vec![empty_hash; tree_size]; + let mut tree = vec![empty_hash; minimum_tree_size]; // Compute and store each leaf hash. tree[num_nodes..num_nodes + leaves.len()].copy_from_slice(&leaf_hasher.hash_leaves(leaves)?); @@ -86,16 +90,33 @@ impl, PH: PathHash // Compute and store the hashes for each level, iterating from the penultimate level to the root level. let mut start_index = num_nodes; + let mut current_empty_node_hash = path_hasher.hash_children(&empty_hash, &empty_hash)?; // Compute the start index of the current level. while let Some(start) = parent(start_index) { // Compute the end index of the current level. let end = left_child(start); // Construct the children for each node in the current level. - let tuples = (start..end).map(|i| (tree[left_child(i)], tree[right_child(i)])).collect::>(); + let tuples = (start..end) + .map(|i| { + // Procure the children of the node. + let tuple = (tree.get(left_child(i)).copied(), tree.get(right_child(i)).copied()); + // If both children are empty hashes, return `None`. + if tuple.0 == Some(current_empty_node_hash) && tuple.1 == Some(current_empty_node_hash) { + return None; + } + // If any of the children are missing, return `None`. + if tuple.0.is_none() || tuple.1.is_none() { + None + } else { + Some((tuple.0.unwrap(), tuple.1.unwrap())) + } + }) + .collect::>(); // Compute and store the hashes for each node in the current level. - tree[start..end].copy_from_slice(&path_hasher.hash_all_children(&tuples)?); + tree[start..end].copy_from_slice(&path_hasher.hash_all_children(&tuples, current_empty_node_hash)?); // Update the start index for the next level. start_index = start; + current_empty_node_hash = path_hasher.hash_children(¤t_empty_node_hash, ¤t_empty_node_hash)?; } lap!(timer, "Hashed {} levels", tree_depth); @@ -144,8 +165,16 @@ impl, PH: PathHash tree.extend(self.leaf_hashes()?); // Extend the new Merkle tree with the new leaf hashes. tree.extend(&self.leaf_hasher.hash_leaves(new_leaves)?); + + // Calculate the size of the tree which excludes leafless nodes. + let new_number_of_leaves = self.number_of_leaves + new_leaves.len(); + let minimum_tree_size = std::cmp::max( + 1, + num_nodes + new_number_of_leaves + if new_number_of_leaves > 1 { new_number_of_leaves % 2 } else { 0 }, + ); + // Resize the new Merkle tree with empty hashes to pad up to `tree_size`. - tree.resize(tree_size, self.empty_hash); + tree.resize(minimum_tree_size, self.empty_hash); lap!(timer, "Hashed {} new leaves", new_leaves.len()); // Initialize a start index to track the starting index of the current level. @@ -453,12 +482,20 @@ impl, PH: PathHash // Compute the number of padded levels. let padding_depth = DEPTH - tree_depth; + // Calculate the size of the tree which excludes leafless nodes. + let minimum_tree_size = std::cmp::max( + 1, + num_nodes + + updated_number_of_leaves + + if updated_number_of_leaves > 1 { updated_number_of_leaves % 2 } else { 0 }, + ); + // Initialize the Merkle tree. let mut tree = vec![self.empty_hash; num_nodes]; // Extend the new Merkle tree with the existing leaf hashes, excluding the last 'n' leaves. tree.extend(&self.leaf_hashes()?[..updated_number_of_leaves]); // Resize the new Merkle tree with empty hashes to pad up to `tree_size`. - tree.resize(tree_size, self.empty_hash); + tree.resize(minimum_tree_size, self.empty_hash); lap!(timer, "Resizing to {} leaves", updated_number_of_leaves); // Initialize a start index to track the starting index of the current level. @@ -627,6 +664,7 @@ impl, PH: PathHash let timer = timer!("MerkleTree::compute_updated_tree"); // Compute and store the hashes for each level, iterating from the penultimate level to the root level. + let empty_hash = self.path_hasher.hash_empty()?; while let (Some(start), Some(middle)) = (parent(start_index), parent(middle_index)) { // Compute the end index of the current level. let end = left_child(start); @@ -651,7 +689,12 @@ impl, PH: PathHash if let Some(middle_precompute) = parent(middle_precompute) { // Construct the children for the new indices in the current level. let tuples = (middle..middle_precompute) - .map(|i| (tree[left_child(i)], tree[right_child(i)])) + .map(|i| { + ( + tree.get(left_child(i)).copied().unwrap_or(empty_hash), + tree.get(right_child(i)).copied().unwrap_or(empty_hash), + ) + }) .collect::>(); // Process the indices that need to be computed for the current level. // If any level requires computing more than 100 nodes, borrow the tree for performance. @@ -687,7 +730,14 @@ impl, PH: PathHash } } else { // Construct the children for the new indices in the current level. - let tuples = (middle..end).map(|i| (tree[left_child(i)], tree[right_child(i)])).collect::>(); + let tuples = (middle..end) + .map(|i| { + ( + tree.get(left_child(i)).copied().unwrap_or(empty_hash), + tree.get(right_child(i)).copied().unwrap_or(empty_hash), + ) + }) + .collect::>(); // Process the indices that need to be computed for the current level. // If any level requires computing more than 100 nodes, borrow the tree for performance. match tuples.len() >= 100 { diff --git a/console/collections/src/merkle_tree/tests/append.rs b/console/collections/src/merkle_tree/tests/append.rs index 968e0b0bd4..41c21102fa 100644 --- a/console/collections/src/merkle_tree/tests/append.rs +++ b/console/collections/src/merkle_tree/tests/append.rs @@ -157,7 +157,7 @@ fn check_merkle_tree_depth_3_padded Date: Thu, 4 Apr 2024 16:14:54 +0200 Subject: [PATCH 2/4] perf: remove redundant checks in MerkleTree::new Signed-off-by: ljedrz --- .../src/merkle_tree/helpers/path_hash.rs | 15 +++--------- console/collections/src/merkle_tree/mod.rs | 24 +++++++++++-------- 2 files changed, 17 insertions(+), 22 deletions(-) diff --git a/console/collections/src/merkle_tree/helpers/path_hash.rs b/console/collections/src/merkle_tree/helpers/path_hash.rs index 83167d0385..fc9841672c 100644 --- a/console/collections/src/merkle_tree/helpers/path_hash.rs +++ b/console/collections/src/merkle_tree/helpers/path_hash.rs @@ -31,20 +31,11 @@ pub trait PathHash: Clone + Send + Sync { fn hash_children(&self, left: &Self::Hash, right: &Self::Hash) -> Result; /// Returns the hash for each tuple of child nodes. - /// If there are no children, the supplied empty hash is used. - fn hash_all_children( - &self, - child_nodes: &[Option<(Self::Hash, Self::Hash)>], - empty_node_hash: Self::Hash, - ) -> Result> { - let hash_children = |children: &Option<(Self::Hash, Self::Hash)>| { - if let Some((left, right)) = children { self.hash_children(left, right) } else { Ok(empty_node_hash) } - }; - + fn hash_all_children(&self, child_nodes: &[(Self::Hash, Self::Hash)]) -> Result> { match child_nodes.len() { 0 => Ok(vec![]), - 1..=100 => child_nodes.iter().map(hash_children).collect(), - _ => cfg_iter!(child_nodes).map(hash_children).collect(), + 1..=100 => child_nodes.iter().map(|(left, right)| self.hash_children(left, right)).collect(), + _ => cfg_iter!(child_nodes).map(|(left, right)| self.hash_children(left, right)).collect(), } } } diff --git a/console/collections/src/merkle_tree/mod.rs b/console/collections/src/merkle_tree/mod.rs index a36e3c65c4..0fe128dede 100644 --- a/console/collections/src/merkle_tree/mod.rs +++ b/console/collections/src/merkle_tree/mod.rs @@ -78,6 +78,8 @@ impl, PH: PathHash let empty_hash = path_hasher.hash_empty()?; // Calculate the size of the tree which excludes leafless nodes. + // The minimum tree size is either a single root node or the calculated number of nodes plus + // the supplied leaves; if the number of leaves is odd, an empty hash is added for padding. let minimum_tree_size = std::cmp::max(1, num_nodes + leaves.len() + if leaves.len() > 1 { leaves.len() % 2 } else { 0 }); @@ -90,22 +92,17 @@ impl, PH: PathHash // Compute and store the hashes for each level, iterating from the penultimate level to the root level. let mut start_index = num_nodes; - let mut current_empty_node_hash = path_hasher.hash_children(&empty_hash, &empty_hash)?; // Compute the start index of the current level. while let Some(start) = parent(start_index) { // Compute the end index of the current level. let end = left_child(start); // Construct the children for each node in the current level. let tuples = (start..end) - .map(|i| { + .filter_map(|i| { // Procure the children of the node. let tuple = (tree.get(left_child(i)).copied(), tree.get(right_child(i)).copied()); - // If both children are empty hashes, return `None`. - if tuple.0 == Some(current_empty_node_hash) && tuple.1 == Some(current_empty_node_hash) { - return None; - } - // If any of the children are missing, return `None`. - if tuple.0.is_none() || tuple.1.is_none() { + // If both children are missing, return `None`. + if tuple.0.is_none() && tuple.1.is_none() { None } else { Some((tuple.0.unwrap(), tuple.1.unwrap())) @@ -113,10 +110,17 @@ impl, PH: PathHash }) .collect::>(); // Compute and store the hashes for each node in the current level. - tree[start..end].copy_from_slice(&path_hasher.hash_all_children(&tuples, current_empty_node_hash)?); + let num_full_nodes = tuples.len(); + tree[start..][..num_full_nodes].copy_from_slice(&path_hasher.hash_all_children(&tuples)?); + // Use the precomputed empty node hash for every empty node, if there are any. + if start + num_full_nodes < end { + let empty_node_hash = path_hasher.hash_children(&empty_hash, &empty_hash)?; + for node in tree.iter_mut().take(end).skip(start + num_full_nodes) { + *node = empty_node_hash; + } + } // Update the start index for the next level. start_index = start; - current_empty_node_hash = path_hasher.hash_children(¤t_empty_node_hash, ¤t_empty_node_hash)?; } lap!(timer, "Hashed {} levels", tree_depth); From 4e2f43472ae9212a4b50e35cf7000d2dc7ac79b7 Mon Sep 17 00:00:00 2001 From: ljedrz Date: Mon, 8 Apr 2024 16:48:20 +0200 Subject: [PATCH 3/4] perf: only check for the presence of a single child in MerkleTree::new Signed-off-by: ljedrz --- console/collections/src/merkle_tree/mod.rs | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/console/collections/src/merkle_tree/mod.rs b/console/collections/src/merkle_tree/mod.rs index 0fe128dede..6c0189c1fd 100644 --- a/console/collections/src/merkle_tree/mod.rs +++ b/console/collections/src/merkle_tree/mod.rs @@ -101,12 +101,8 @@ impl, PH: PathHash .filter_map(|i| { // Procure the children of the node. let tuple = (tree.get(left_child(i)).copied(), tree.get(right_child(i)).copied()); - // If both children are missing, return `None`. - if tuple.0.is_none() && tuple.1.is_none() { - None - } else { - Some((tuple.0.unwrap(), tuple.1.unwrap())) - } + // If the left child is missing (in which case the right one also is), return `None`. + if tuple.0.is_none() { None } else { Some((tuple.0.unwrap(), tuple.1.unwrap())) } }) .collect::>(); // Compute and store the hashes for each node in the current level. From 7fc8a2b16e3099b9ae709913cf5558863105d653 Mon Sep 17 00:00:00 2001 From: ljedrz Date: Mon, 8 Apr 2024 16:58:42 +0200 Subject: [PATCH 4/4] perf: stop iterating the leaves sooner in MerkleTree::new Signed-off-by: ljedrz --- console/collections/src/merkle_tree/mod.rs | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/console/collections/src/merkle_tree/mod.rs b/console/collections/src/merkle_tree/mod.rs index 6c0189c1fd..cf3d2aeb61 100644 --- a/console/collections/src/merkle_tree/mod.rs +++ b/console/collections/src/merkle_tree/mod.rs @@ -96,14 +96,11 @@ impl, PH: PathHash while let Some(start) = parent(start_index) { // Compute the end index of the current level. let end = left_child(start); - // Construct the children for each node in the current level. + // Construct the children for each node in the current level; the leaves are padded, which means + // that there either are 2 children, or there are none, at which point we may stop iterating. let tuples = (start..end) - .filter_map(|i| { - // Procure the children of the node. - let tuple = (tree.get(left_child(i)).copied(), tree.get(right_child(i)).copied()); - // If the left child is missing (in which case the right one also is), return `None`. - if tuple.0.is_none() { None } else { Some((tuple.0.unwrap(), tuple.1.unwrap())) } - }) + .take_while(|&i| tree.get(left_child(i)).is_some()) + .map(|i| (tree[left_child(i)], tree[right_child(i)])) .collect::>(); // Compute and store the hashes for each node in the current level. let num_full_nodes = tuples.len();