trivial: less vector allocation

aptos-labs · May 13, 2024 · fa2fa03 · fa2fa03
1 parent 4eae1fa
commit fa2fa03
Show file tree

Hide file tree

Showing 3 changed files with 82 additions and 66 deletions.
diff --git a/storage/jellyfish-merkle/src/lib.rs b/storage/jellyfish-merkle/src/lib.rs
@@ -711,7 +711,7 @@ where
     ) -> Result<(Option<(HashValue, (K, Version))>, SparseMerkleProofExt)> {
         // Empty tree just returns proof with no sibling hash.
         let mut next_node_key = NodeKey::new_empty_path(version);
-        let mut siblings = vec![];
+        let mut out_siblings = vec![];
         let nibble_path = NibblePath::new_even(key.to_vec());
         let mut nibble_iter = nibble_path.nibbles();
 
@@ -742,21 +742,20 @@ where
                     let queried_child_index = nibble_iter
                         .next()
                         .ok_or_else(|| AptosDbError::Other("ran out of nibbles".to_string()))?;
-                    let (child_node_key, mut siblings_in_internal) = internal_node
-                        .get_child_with_siblings(
-                            &next_node_key,
-                            queried_child_index,
-                            Some(self.reader),
-                        )?;
-                    siblings.append(&mut siblings_in_internal);
+                    let child_node_key = internal_node.get_child_with_siblings(
+                        &next_node_key,
+                        queried_child_index,
+                        Some(self.reader),
+                        &mut out_siblings,
+                    )?;
                     next_node_key = match child_node_key {
                         Some(node_key) => node_key,
                         None => {
                             return Ok((
                                 None,
                                 SparseMerkleProofExt::new(None, {
-                                    siblings.reverse();
-                                    siblings
+                                    out_siblings.reverse();
+                                    out_siblings
                                 }),
                             ));
                         },
@@ -770,8 +769,8 @@ where
                             None
                         },
                         SparseMerkleProofExt::new(Some(leaf_node.into()), {
-                            siblings.reverse();
-                            siblings
+                            out_siblings.reverse();
+                            out_siblings
                         }),
                     ));
                 },

diff --git a/storage/jellyfish-merkle/src/node_type/mod.rs b/storage/jellyfish-merkle/src/node_type/mod.rs
@@ -559,10 +559,10 @@ impl InternalNode {
         node_key: &NodeKey,
         n: Nibble,
         reader: Option<&R>,
-    ) -> Result<(Option<NodeKey>, Vec<NodeInProof>)> {
+        out_siblings: &mut Vec<NodeInProof>,
+    ) -> Result<Option<NodeKey>> {
         assert!(self.leaf_count > 1);
 
-        let mut siblings = vec![];
         let (existence_bitmap, leaf_bitmap) = self.generate_bitmaps();
 
         // Nibble height from 3 to 0.
@@ -573,14 +573,14 @@ impl InternalNode {
             let (child_half_start, sibling_half_start) = get_child_and_sibling_half_start(n, h);
             // Compute the root hash of the subtree rooted at the sibling of `r`.
             if let Some(reader) = reader {
-                siblings.push(self.gen_node_in_proof(
+                out_siblings.push(self.gen_node_in_proof(
                     sibling_half_start,
                     width,
                     (existence_bitmap, leaf_bitmap),
                     (reader, node_key),
                 )?);
             } else {
-                siblings.push(
+                out_siblings.push(
                     self.merkle_hash(sibling_half_start, width, (existence_bitmap, leaf_bitmap))
                         .into(),
                 );
@@ -591,7 +591,7 @@ impl InternalNode {
 
             if range_existence_bitmap == 0 {
                 // No child in this range.
-                return Ok((None, siblings));
+                return Ok(None);
             } else if width == 1
                 || (range_existence_bitmap.count_ones() == 1 && range_leaf_bitmap != 0)
             {
@@ -600,28 +600,37 @@ impl InternalNode {
                 // `None` because it's existence indirectly proves the n-th child doesn't exist.
                 // Please read proof format for details.
                 let only_child_index = Nibble::from(range_existence_bitmap.trailing_zeros() as u8);
-                return Ok((
-                    {
-                        let only_child_version = self
-                            .child(only_child_index)
-                            // Should be guaranteed by the self invariants, but these are not easy to express at the moment
-                            .with_context(|| {
-                                format!(
-                                    "Corrupted internal node: child_bitmap indicates \
-                                     the existence of a non-exist child at index {:x}",
-                                    only_child_index
-                                )
-                            })
-                            .unwrap()
-                            .version;
-                        Some(node_key.gen_child_node_key(only_child_version, only_child_index))
-                    },
-                    siblings,
+                let only_child_version = self
+                    .child(only_child_index)
+                    // Should be guaranteed by the self invariants, but these are not easy to express at the moment
+                    .with_context(|| {
+                        format!(
+                            "Corrupted internal node: child_bitmap indicates \
+                                 the existence of a non-exist child at index {:x}",
+                            only_child_index
+                        )
+                    })
+                    .unwrap()
+                    .version;
+                return Ok(Some(
+                    node_key.gen_child_node_key(only_child_version, only_child_index),
                 ));
             }
         }
         unreachable!("Impossible to get here without returning even at the lowest level.")
     }
+
+    #[cfg(test)]
+    pub(crate) fn get_child_with_siblings_for_test<K: crate::Key, R: TreeReader<K>>(
+        &self,
+        node_key: &NodeKey,
+        n: Nibble,
+        reader: Option<&R>,
+    ) -> Result<(Option<NodeKey>, Vec<NodeInProof>)> {
+        let mut sibilings = vec![];
+        self.get_child_with_siblings(node_key, n, reader, &mut sibilings)
+            .map(|n| (n, sibilings))
+    }
 }
 
 /// Given a nibble, computes the start position of its `child_half_start` and `sibling_half_start`