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..85ca8f5924 100644 --- a/console/collections/src/kary_merkle_tree/helpers/path_hash.rs +++ b/console/collections/src/kary_merkle_tree/helpers/path_hash.rs @@ -32,8 +32,8 @@ 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. + 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 81ebd5c488..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> { @@ -89,8 +90,18 @@ 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. + // 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 all_nodes_are_full { 0 } else { arity - leaves.len() % 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)?); @@ -105,10 +116,21 @@ 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::>()) + .take_while(|&i| child_indexes::(i).next().and_then(|idx| tree.get(idx)).is_some()) + .map(|i| &tree[child_indexes::(i)]) .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)?); + let num_full_nodes = child_nodes.len(); + 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])?; + 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; } @@ -241,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 } 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]); }