Reduce zero-padding and cache zero-hashes in k-ary MerkleTree

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<Self::Hash>]) -> Result<Vec<Self::Hash>> {
+    /// 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<Vec<Self::Hash>>],
+        empty_node_hash: Self::Hash,
+    ) -> Result<Vec<Self::Hash>> {
+        let hash_children = |children: &Option<Vec<Self::Hash>>| {
+            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(),
         }
     }
 }

console/collections/src/kary_merkle_tree/mod.rs

@@ -89,28 +89,55 @@ impl<LH: LeafHash<Hash = PH::Hash>, PH: PathHash, const DEPTH: u8, const ARITY:
         // Compute the empty hash.
         let empty_hash = path_hasher.hash_empty::<ARITY>()?;
 
+        // 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 },
+        );
+
         // 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)?);
         lap!(timer, "Hashed {} leaves", leaves.len());
 
         // 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::<ARITY>(start_index) {
             // Compute the end index of the current level.
             let end = child_indexes::<ARITY>(start).next().ok_or_else(|| anyhow!("Missing left-most child"))?;
 
             // Construct the children for each node in the current level.
             let child_nodes = (start..end)
-                .map(|i| child_indexes::<ARITY>(i).map(|child_index| tree[child_index]).collect::<Vec<_>>())
+                .map(|i| {
+                    // Prepare an iterator over then children, being mindful of possible missing leaves.
+                    let child_iter = || child_indexes::<ARITY>(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::<Option<Vec<_>>>();
+                    }
+
+                    // 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::<Option<Vec<_>>>()
+                })
                 .collect::<Vec<_>>();
+
             // 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, current_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);
 

console/collections/src/kary_merkle_tree/tests/mod.rs

@@ -170,7 +170,7 @@ fn check_merkle_tree_depth_3_arity_3_padded<LH: LeafHash<Hash = PH::Hash>, PH: P
 
     // Construct the Merkle tree for the given leaves.
     let merkle_tree = KaryMerkleTree::<LH, PH, 3, ARITY>::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<LH: LeafHash<Hash = PH::Hash>, 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::<ARITY>()?, merkle_tree.tree[i]);
     }