From 362772065f5087d653efce4ed0c918d7b6719b67 Mon Sep 17 00:00:00 2001 From: ljedrz Date: Fri, 22 Mar 2024 11:03:28 +0100 Subject: [PATCH 1/9] perf: reduce zero-padding and cache zero-hashes in k-ary MerkleTree Signed-off-by: ljedrz --- .../src/kary_merkle_tree/helpers/path_hash.rs | 17 ++++++-- .../collections/src/kary_merkle_tree/mod.rs | 42 +++++++++++++++++-- .../src/kary_merkle_tree/tests/mod.rs | 4 +- 3 files changed, 54 insertions(+), 9 deletions(-) diff --git a/console/collections/src/kary_merkle_tree/helpers/path_hash.rs b/console/collections/src/kary_merkle_tree/helpers/path_hash.rs index 57ae61a2a4..22146191ef 100644 --- a/console/collections/src/kary_merkle_tree/helpers/path_hash.rs +++ b/console/collections/src/kary_merkle_tree/helpers/path_hash.rs @@ -32,12 +32,21 @@ pub trait PathHash: Clone + Send + Sync { self.hash_children(&children) } - /// Returns the hash for each tuple of child nodes. - fn hash_all_children(&self, child_nodes: &[Vec]) -> Result> { + /// Returns the hash for each child node. + /// If there are no children, the supplied empty node hash is used instead. + fn hash_all_children( + &self, + child_nodes: &[Option>], + empty_node_hash: Self::Hash, + ) -> Result> { + let hash_children = |children: &Option>| { + if let Some(children) = children { self.hash_children(children) } else { Ok(empty_node_hash) } + }; + match child_nodes.len() { 0 => Ok(vec![]), - 1..=100 => child_nodes.iter().map(|children| self.hash_children(children)).collect(), - _ => cfg_iter!(child_nodes).map(|children| self.hash_children(children)).collect(), + 1..=100 => child_nodes.iter().map(hash_children).collect(), + _ => cfg_iter!(child_nodes).map(hash_children).collect(), } } } diff --git a/console/collections/src/kary_merkle_tree/mod.rs b/console/collections/src/kary_merkle_tree/mod.rs index 81ebd5c488..abf8922f6d 100644 --- a/console/collections/src/kary_merkle_tree/mod.rs +++ b/console/collections/src/kary_merkle_tree/mod.rs @@ -89,8 +89,22 @@ impl, PH: PathHash, const DEPTH: u8, const ARITY: // Compute the empty hash. let empty_hash = path_hasher.hash_empty::()?; + // Calculate the size of the tree which excludes leafless nodes. + let arity = ARITY as usize; + let minimum_tree_size = std::cmp::max( + 1, + num_nodes + leaves.len() + if leaves.len() % arity == 0 { 0 } else { arity - leaves.len() % arity }, + ); + + // Prepare a cache of hashes of empty nodes and their derivatives. + let mut empty_hash_cache = Vec::with_capacity(DEPTH as usize); + empty_hash_cache.push(path_hasher.hash_children(&vec![empty_hash; arity])?); + for i in 0..(DEPTH as usize - 1) { + empty_hash_cache.push(path_hasher.hash_children(&vec![empty_hash_cache[i]; arity])?); + } + // 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()].clone_from_slice(&leaf_hasher.hash_leaves(leaves)?); @@ -98,6 +112,7 @@ impl, PH: PathHash, const DEPTH: u8, const ARITY: // 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_level = 0; // Compute the start index of the current level. while let Some(start) = parent::(start_index) { // Compute the end index of the current level. @@ -105,12 +120,33 @@ impl, PH: PathHash, const DEPTH: u8, const ARITY: // Construct the children for each node in the current level. let child_nodes = (start..end) - .map(|i| child_indexes::(i).map(|child_index| tree[child_index]).collect::>()) + .map(|i| { + // Prepare an iterator over then children, being mindful of possible missing leaves. + let child_iter = || child_indexes::(i).map(|child_index| tree.get(child_index).copied()); + + // At the leaf level just return the leaves, or `None` in case the node has none. + if current_level == 0 { + return child_iter().collect::>>(); + } + + // Check if the children aren't all derived from empty hashes. + let current_composite_empty_hash = empty_hash_cache[current_level - 1]; + if child_iter().all(|hash| hash == Some(current_composite_empty_hash)) { + return None; + } + + // Collect the children. + child_iter().collect::>>() + }) .collect::>(); + // Compute and store the hashes for each node in the current level. - tree[start..end].clone_from_slice(&path_hasher.hash_all_children(&child_nodes)?); + tree[start..end].clone_from_slice( + &path_hasher.hash_all_children::(&child_nodes, empty_hash_cache[current_level])?, + ); // Update the start index for the next level. start_index = start; + current_level += 1; } lap!(timer, "Hashed {} levels", tree_depth); diff --git a/console/collections/src/kary_merkle_tree/tests/mod.rs b/console/collections/src/kary_merkle_tree/tests/mod.rs index 08abc90673..10a99f6a0f 100644 --- a/console/collections/src/kary_merkle_tree/tests/mod.rs +++ b/console/collections/src/kary_merkle_tree/tests/mod.rs @@ -170,7 +170,7 @@ fn check_merkle_tree_depth_3_arity_3_padded, PH: P // Construct the Merkle tree for the given leaves. let merkle_tree = KaryMerkleTree::::new(leaf_hasher, path_hasher, &leaves)?; - assert_eq!(40, merkle_tree.tree.len()); + assert_eq!(25, merkle_tree.tree.len()); assert_eq!(10, merkle_tree.number_of_leaves); // Depth 3. @@ -194,7 +194,7 @@ fn check_merkle_tree_depth_3_arity_3_padded, PH: P assert_eq!(expected_leaf7, merkle_tree.tree[20]); assert_eq!(expected_leaf8, merkle_tree.tree[21]); assert_eq!(expected_leaf9, merkle_tree.tree[22]); - for i in 23..40 { + for i in 23..25 { assert_eq!(path_hasher.hash_empty::()?, merkle_tree.tree[i]); } From 3afcadd791c84af160a53d02f4cc3587f367cdd0 Mon Sep 17 00:00:00 2001 From: ljedrz Date: Sat, 23 Mar 2024 17:38:22 +0100 Subject: [PATCH 2/9] perf: only cache the zero hash for the previous level Signed-off-by: ljedrz --- .../src/kary_merkle_tree/helpers/path_hash.rs | 2 +- .../collections/src/kary_merkle_tree/mod.rs | 21 ++++++------------- 2 files changed, 7 insertions(+), 16 deletions(-) diff --git a/console/collections/src/kary_merkle_tree/helpers/path_hash.rs b/console/collections/src/kary_merkle_tree/helpers/path_hash.rs index 22146191ef..6087e4800e 100644 --- a/console/collections/src/kary_merkle_tree/helpers/path_hash.rs +++ b/console/collections/src/kary_merkle_tree/helpers/path_hash.rs @@ -34,7 +34,7 @@ pub trait PathHash: Clone + Send + Sync { /// Returns the hash for each child node. /// If there are no children, the supplied empty node hash is used instead. - fn hash_all_children( + fn hash_all_children( &self, child_nodes: &[Option>], empty_node_hash: Self::Hash, diff --git a/console/collections/src/kary_merkle_tree/mod.rs b/console/collections/src/kary_merkle_tree/mod.rs index abf8922f6d..65fefd1a4e 100644 --- a/console/collections/src/kary_merkle_tree/mod.rs +++ b/console/collections/src/kary_merkle_tree/mod.rs @@ -96,13 +96,6 @@ impl, PH: PathHash, const DEPTH: u8, const ARITY: num_nodes + leaves.len() + if leaves.len() % arity == 0 { 0 } else { arity - leaves.len() % arity }, ); - // Prepare a cache of hashes of empty nodes and their derivatives. - let mut empty_hash_cache = Vec::with_capacity(DEPTH as usize); - empty_hash_cache.push(path_hasher.hash_children(&vec![empty_hash; arity])?); - for i in 0..(DEPTH as usize - 1) { - empty_hash_cache.push(path_hasher.hash_children(&vec![empty_hash_cache[i]; arity])?); - } - // Initialize the Merkle tree. let mut tree = vec![empty_hash; minimum_tree_size]; @@ -112,7 +105,7 @@ impl, PH: PathHash, const DEPTH: u8, const ARITY: // 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_level = 0; + let mut current_empty_node_hash = path_hasher.hash_children(&vec![empty_hash; arity])?; // Compute the start index of the current level. while let Some(start) = parent::(start_index) { // Compute the end index of the current level. @@ -125,13 +118,12 @@ impl, PH: PathHash, const DEPTH: u8, const ARITY: let child_iter = || child_indexes::(i).map(|child_index| tree.get(child_index).copied()); // At the leaf level just return the leaves, or `None` in case the node has none. - if current_level == 0 { + if current_empty_node_hash == empty_hash { return child_iter().collect::>>(); } // Check if the children aren't all derived from empty hashes. - let current_composite_empty_hash = empty_hash_cache[current_level - 1]; - if child_iter().all(|hash| hash == Some(current_composite_empty_hash)) { + if child_iter().all(|hash| hash == Some(current_empty_node_hash)) { return None; } @@ -141,12 +133,11 @@ impl, PH: PathHash, const DEPTH: u8, const ARITY: .collect::>(); // Compute and store the hashes for each node in the current level. - tree[start..end].clone_from_slice( - &path_hasher.hash_all_children::(&child_nodes, empty_hash_cache[current_level])?, - ); + tree[start..end].clone_from_slice(&path_hasher.hash_all_children(&child_nodes, current_empty_node_hash)?); // Update the start index for the next level. start_index = start; - current_level += 1; + // Update the empty node hash for the next level. + current_empty_node_hash = path_hasher.hash_children(&vec![current_empty_node_hash; arity])?; } lap!(timer, "Hashed {} levels", tree_depth); From fb429a4633e9d06f2ddabf407e218e5e13e571d2 Mon Sep 17 00:00:00 2001 From: ljedrz Date: Thu, 4 Apr 2024 14:44:15 +0200 Subject: [PATCH 3/9] docs: add an extra comment in KaryMerkleTree::new Signed-off-by: ljedrz --- console/collections/src/kary_merkle_tree/mod.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/console/collections/src/kary_merkle_tree/mod.rs b/console/collections/src/kary_merkle_tree/mod.rs index 65fefd1a4e..d4d3636095 100644 --- a/console/collections/src/kary_merkle_tree/mod.rs +++ b/console/collections/src/kary_merkle_tree/mod.rs @@ -90,10 +90,13 @@ impl, PH: PathHash, const DEPTH: u8, const ARITY: 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, and empty hashes that pad up to the tree's arity (making every node full). let arity = ARITY as usize; + let all_nodes_are_full = leaves.len() % arity == 0; let minimum_tree_size = std::cmp::max( 1, - num_nodes + leaves.len() + if leaves.len() % arity == 0 { 0 } else { arity - leaves.len() % arity }, + num_nodes + leaves.len() + if all_nodes_are_full { 0 } else { arity - leaves.len() % arity }, ); // Initialize the Merkle tree. From 6bc448ce3622527182c0e86f67e6244623f053cd Mon Sep 17 00:00:00 2001 From: ljedrz Date: Thu, 4 Apr 2024 14:59:59 +0200 Subject: [PATCH 4/9] cleanup: remove a redundant check in KaryMerkleTree::new Signed-off-by: ljedrz --- console/collections/src/kary_merkle_tree/mod.rs | 5 ----- 1 file changed, 5 deletions(-) diff --git a/console/collections/src/kary_merkle_tree/mod.rs b/console/collections/src/kary_merkle_tree/mod.rs index d4d3636095..cd272927e7 100644 --- a/console/collections/src/kary_merkle_tree/mod.rs +++ b/console/collections/src/kary_merkle_tree/mod.rs @@ -120,11 +120,6 @@ impl, PH: PathHash, const DEPTH: u8, const ARITY: // Prepare an iterator over then children, being mindful of possible missing leaves. let child_iter = || child_indexes::(i).map(|child_index| tree.get(child_index).copied()); - // At the leaf level just return the leaves, or `None` in case the node has none. - if current_empty_node_hash == empty_hash { - return child_iter().collect::>>(); - } - // Check if the children aren't all derived from empty hashes. if child_iter().all(|hash| hash == Some(current_empty_node_hash)) { return None; From cd2fe47fd703ee4f11183d4f8a482b9b3f84f743 Mon Sep 17 00:00:00 2001 From: ljedrz Date: Thu, 4 Apr 2024 15:12:39 +0200 Subject: [PATCH 5/9] cleanup: remove another redundant check in KaryMerkleTree::new Signed-off-by: ljedrz --- .../src/kary_merkle_tree/helpers/path_hash.rs | 14 ++------- .../collections/src/kary_merkle_tree/mod.rs | 29 +++++++++---------- 2 files changed, 17 insertions(+), 26 deletions(-) diff --git a/console/collections/src/kary_merkle_tree/helpers/path_hash.rs b/console/collections/src/kary_merkle_tree/helpers/path_hash.rs index 6087e4800e..287bc789d1 100644 --- a/console/collections/src/kary_merkle_tree/helpers/path_hash.rs +++ b/console/collections/src/kary_merkle_tree/helpers/path_hash.rs @@ -34,19 +34,11 @@ pub trait PathHash: Clone + Send + Sync { /// Returns the hash for each child node. /// If there are no children, the supplied empty node hash is used instead. - fn hash_all_children( - &self, - child_nodes: &[Option>], - empty_node_hash: Self::Hash, - ) -> Result> { - let hash_children = |children: &Option>| { - if let Some(children) = children { self.hash_children(children) } else { Ok(empty_node_hash) } - }; - + fn hash_all_children(&self, child_nodes: &[Vec]) -> 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(|children| self.hash_children(children)).collect(), + _ => cfg_iter!(child_nodes).map(|children| self.hash_children(children)).collect(), } } } diff --git a/console/collections/src/kary_merkle_tree/mod.rs b/console/collections/src/kary_merkle_tree/mod.rs index cd272927e7..daa3d6a63b 100644 --- a/console/collections/src/kary_merkle_tree/mod.rs +++ b/console/collections/src/kary_merkle_tree/mod.rs @@ -108,7 +108,6 @@ impl, PH: PathHash, const DEPTH: u8, const ARITY: // 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(&vec![empty_hash; arity])?; // Compute the start index of the current level. while let Some(start) = parent::(start_index) { // Compute the end index of the current level. @@ -116,26 +115,26 @@ impl, PH: PathHash, const DEPTH: u8, const ARITY: // Construct the children for each node in the current level. let child_nodes = (start..end) - .map(|i| { - // Prepare an iterator over then children, being mindful of possible missing leaves. - let child_iter = || child_indexes::(i).map(|child_index| tree.get(child_index).copied()); - - // Check if the children aren't all derived from empty hashes. - if child_iter().all(|hash| hash == Some(current_empty_node_hash)) { - return None; - } - - // Collect the children. - child_iter().collect::>>() + .filter_map(|i| { + // Collect the children, being mindful of possible missing leaves. + child_indexes::(i) + .map(|child_index| tree.get(child_index).copied()) + .collect::>>() }) .collect::>(); // Compute and store the hashes for each node in the current level. - tree[start..end].clone_from_slice(&path_hasher.hash_all_children(&child_nodes, current_empty_node_hash)?); + let num_full_nodes = child_nodes.len(); + tree[start..][..num_full_nodes].clone_from_slice(&path_hasher.hash_all_children(&child_nodes)?); + // 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(&vec![empty_hash; arity])?; + for idx in start + num_full_nodes..end { + tree[idx] = empty_node_hash; + } + } // Update the start index for the next level. start_index = start; - // Update the empty node hash for the next level. - current_empty_node_hash = path_hasher.hash_children(&vec![current_empty_node_hash; arity])?; } lap!(timer, "Hashed {} levels", tree_depth); From 7646a98beec235f9177bd6d0820e56f14adf1ad9 Mon Sep 17 00:00:00 2001 From: ljedrz Date: Thu, 4 Apr 2024 15:57:30 +0200 Subject: [PATCH 6/9] docs: fix the doc comment for hash_all_children Signed-off-by: ljedrz --- console/collections/src/kary_merkle_tree/helpers/path_hash.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/console/collections/src/kary_merkle_tree/helpers/path_hash.rs b/console/collections/src/kary_merkle_tree/helpers/path_hash.rs index 287bc789d1..de41116d70 100644 --- a/console/collections/src/kary_merkle_tree/helpers/path_hash.rs +++ b/console/collections/src/kary_merkle_tree/helpers/path_hash.rs @@ -33,7 +33,6 @@ pub trait PathHash: Clone + Send + Sync { } /// Returns the hash for each child node. - /// If there are no children, the supplied empty node hash is used instead. fn hash_all_children(&self, child_nodes: &[Vec]) -> Result> { match child_nodes.len() { 0 => Ok(vec![]), From 650271daef597da1c0330b83b5df27190a2c6419 Mon Sep 17 00:00:00 2001 From: ljedrz Date: Thu, 4 Apr 2024 16:05:47 +0200 Subject: [PATCH 7/9] clippy: improve iteration Signed-off-by: ljedrz --- console/collections/src/kary_merkle_tree/mod.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/console/collections/src/kary_merkle_tree/mod.rs b/console/collections/src/kary_merkle_tree/mod.rs index daa3d6a63b..d37f9254b6 100644 --- a/console/collections/src/kary_merkle_tree/mod.rs +++ b/console/collections/src/kary_merkle_tree/mod.rs @@ -129,8 +129,8 @@ impl, PH: PathHash, const DEPTH: u8, const ARITY: // 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(&vec![empty_hash; arity])?; - for idx in start + num_full_nodes..end { - tree[idx] = empty_node_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. From 2d2b8106ed9c80cd5e48c7122ffef7c75ad60e0a Mon Sep 17 00:00:00 2001 From: ljedrz Date: Mon, 8 Apr 2024 17:06:10 +0200 Subject: [PATCH 8/9] perf: stop iterating the leaves sooner in KaryMerkleTree::new Signed-off-by: ljedrz --- console/collections/src/kary_merkle_tree/mod.rs | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/console/collections/src/kary_merkle_tree/mod.rs b/console/collections/src/kary_merkle_tree/mod.rs index d37f9254b6..87757e4c08 100644 --- a/console/collections/src/kary_merkle_tree/mod.rs +++ b/console/collections/src/kary_merkle_tree/mod.rs @@ -115,12 +115,8 @@ impl, PH: PathHash, const DEPTH: u8, const ARITY: // Construct the children for each node in the current level. let child_nodes = (start..end) - .filter_map(|i| { - // Collect the children, being mindful of possible missing leaves. - child_indexes::(i) - .map(|child_index| tree.get(child_index).copied()) - .collect::>>() - }) + .take_while(|&i| child_indexes::(i).next().and_then(|idx| tree.get(idx)).is_some()) + .map(|i| child_indexes::(i).map(|child_index| tree[child_index]).collect::>()) .collect::>(); // Compute and store the hashes for each node in the current level. From e7031f46d7ec383d339b6789213f5f895546ca83 Mon Sep 17 00:00:00 2001 From: ljedrz Date: Mon, 8 Apr 2024 17:13:55 +0200 Subject: [PATCH 9/9] perf: don't allocate the children in KaryMerkleTree::now Signed-off-by: ljedrz --- .../collections/src/kary_merkle_tree/helpers/path_hash.rs | 2 +- console/collections/src/kary_merkle_tree/mod.rs | 8 +++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/console/collections/src/kary_merkle_tree/helpers/path_hash.rs b/console/collections/src/kary_merkle_tree/helpers/path_hash.rs index de41116d70..85ca8f5924 100644 --- a/console/collections/src/kary_merkle_tree/helpers/path_hash.rs +++ b/console/collections/src/kary_merkle_tree/helpers/path_hash.rs @@ -33,7 +33,7 @@ pub trait PathHash: Clone + Send + Sync { } /// Returns the hash for each child node. - fn hash_all_children(&self, child_nodes: &[Vec]) -> Result> { + fn hash_all_children(&self, child_nodes: &[&[Self::Hash]]) -> Result> { match child_nodes.len() { 0 => Ok(vec![]), 1..=100 => child_nodes.iter().map(|children| self.hash_children(children)).collect(), diff --git a/console/collections/src/kary_merkle_tree/mod.rs b/console/collections/src/kary_merkle_tree/mod.rs index 87757e4c08..3ee7eaf113 100644 --- a/console/collections/src/kary_merkle_tree/mod.rs +++ b/console/collections/src/kary_merkle_tree/mod.rs @@ -24,6 +24,7 @@ mod tests; use snarkvm_console_types::prelude::*; use aleo_std::prelude::*; +use std::ops::Range; #[derive(Clone)] pub struct KaryMerkleTree, PH: PathHash, const DEPTH: u8, const ARITY: u8> { @@ -116,12 +117,13 @@ impl, PH: PathHash, const DEPTH: u8, const ARITY: // Construct the children for each node in the current level. let child_nodes = (start..end) .take_while(|&i| child_indexes::(i).next().and_then(|idx| tree.get(idx)).is_some()) - .map(|i| child_indexes::(i).map(|child_index| tree[child_index]).collect::>()) + .map(|i| &tree[child_indexes::(i)]) .collect::>(); // Compute and store the hashes for each node in the current level. let num_full_nodes = child_nodes.len(); - tree[start..][..num_full_nodes].clone_from_slice(&path_hasher.hash_all_children(&child_nodes)?); + let hashes = path_hasher.hash_all_children(&child_nodes)?; + tree[start..][..num_full_nodes].clone_from_slice(&hashes); // 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(&vec![empty_hash; arity])?; @@ -261,7 +263,7 @@ fn tree_depth(tree_size: usize) -> Result } /// Returns the indexes of the children, given an index. -fn child_indexes(index: usize) -> impl Iterator { +fn child_indexes(index: usize) -> Range { let start = index * ARITY as usize + 1; start..start + ARITY as usize }