From 649b8c64b96282b747bcaec77e7515f57e1fc1bf Mon Sep 17 00:00:00 2001 From: aldenhu Date: Mon, 13 May 2024 22:06:26 +0000 Subject: [PATCH] trivial: less vector allocation --- storage/jellyfish-merkle/src/lib.rs | 23 +++--- storage/jellyfish-merkle/src/node_type/mod.rs | 53 ++++++++------ .../src/node_type/node_type_test.rs | 72 ++++++++++--------- 3 files changed, 82 insertions(+), 66 deletions(-) diff --git a/storage/jellyfish-merkle/src/lib.rs b/storage/jellyfish-merkle/src/lib.rs index 6418e61a4127f..d7fbf2cfe2f7d 100644 --- a/storage/jellyfish-merkle/src/lib.rs +++ b/storage/jellyfish-merkle/src/lib.rs @@ -728,7 +728,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(); @@ -759,21 +759,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 }), )); }, @@ -787,8 +786,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 index 26b1740fb9510..e2f7048f8f98a 100644 --- a/storage/jellyfish-merkle/src/node_type/mod.rs +++ b/storage/jellyfish-merkle/src/node_type/mod.rs @@ -597,10 +597,10 @@ impl InternalNode { node_key: &NodeKey, n: Nibble, reader: Option<&R>, - ) -> Result<(Option, Vec)> { + out_siblings: &mut Vec, + ) -> Result> { assert!(self.leaf_count > 1); - let mut siblings = vec![]; let (existence_bitmap, leaf_bitmap) = self.generate_bitmaps(); // Nibble height from 3 to 0. @@ -611,14 +611,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(), ); @@ -629,7 +629,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) { @@ -638,28 +638,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>( + &self, + node_key: &NodeKey, + n: Nibble, + reader: Option<&R>, + ) -> Result<(Option, Vec)> { + 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` diff --git a/storage/jellyfish-merkle/src/node_type/node_type_test.rs b/storage/jellyfish-merkle/src/node_type/node_type_test.rs index ab1f7e1f21439..77cab6a7f5df0 100644 --- a/storage/jellyfish-merkle/src/node_type/node_type_test.rs +++ b/storage/jellyfish-merkle/src/node_type/node_type_test.rs @@ -196,13 +196,13 @@ proptest! { for i in 0..8 { prop_assert_eq!( - internal_node.get_child_with_siblings::(&internal_node_key, i.into(), None).unwrap(), + internal_node.get_child_with_siblings_for_test::(&internal_node_key, i.into(), None).unwrap(), (Some(leaf1_node_key.clone()), vec![hash2.into()]) ); } for i in 8..16 { prop_assert_eq!( - internal_node.get_child_with_siblings::(&internal_node_key, i.into(), None).unwrap(), + internal_node.get_child_with_siblings_for_test::(&internal_node_key, i.into(), None).unwrap(), (Some(leaf2_node_key.clone()), vec![hash1.into()]) ); } @@ -243,14 +243,14 @@ proptest! { for i in 0..4 { prop_assert_eq!( - internal_node.get_child_with_siblings::(&internal_node_key, i.into(), None).unwrap(), + internal_node.get_child_with_siblings_for_test::(&internal_node_key, i.into(), None).unwrap(), (None, vec![(*SPARSE_MERKLE_PLACEHOLDER_HASH).into(), hash_x1.into()]) ); } for i in 4..6 { prop_assert_eq!( - internal_node.get_child_with_siblings::(&internal_node_key, i.into(), None).unwrap(), + internal_node.get_child_with_siblings_for_test::(&internal_node_key, i.into(), None).unwrap(), ( Some(leaf1_node_key.clone()), vec![ @@ -264,7 +264,7 @@ proptest! { for i in 6..8 { prop_assert_eq!( - internal_node.get_child_with_siblings::(&internal_node_key, i.into(), None).unwrap(), + internal_node.get_child_with_siblings_for_test::(&internal_node_key, i.into(), None).unwrap(), ( Some(leaf2_node_key.clone()), vec![ @@ -278,7 +278,7 @@ proptest! { for i in 8..16 { prop_assert_eq!( - internal_node.get_child_with_siblings::(&internal_node_key, i.into(), None).unwrap(), + internal_node.get_child_with_siblings_for_test::(&internal_node_key, i.into(), None).unwrap(), (None, vec![hash_x2.into()]) ); } @@ -317,21 +317,21 @@ proptest! { for i in 0..4 { prop_assert_eq!( - internal_node.get_child_with_siblings::(&internal_node_key, i.into(), None).unwrap(), + internal_node.get_child_with_siblings_for_test::(&internal_node_key, i.into(), None).unwrap(), (Some(leaf1_node_key.clone()),vec![hash3.into(), hash2.into()]) ); } for i in 4..8 { prop_assert_eq!( - internal_node.get_child_with_siblings::(&internal_node_key, i.into(), None).unwrap(), + internal_node.get_child_with_siblings_for_test::(&internal_node_key, i.into(), None).unwrap(), (Some(leaf2_node_key.clone()),vec![hash3.into(), hash1.into()]) ); } for i in 8..16 { prop_assert_eq!( - internal_node.get_child_with_siblings::(&internal_node_key, i.into(), None).unwrap(), + internal_node.get_child_with_siblings_for_test::(&internal_node_key, i.into(), None).unwrap(), (Some(leaf3_node_key.clone()),vec![hash_x.into()]) ); } @@ -382,7 +382,7 @@ proptest! { for i in 0..2 { prop_assert_eq!( - internal_node.get_child_with_siblings::(&internal_node_key, i.into(), None).unwrap(), + internal_node.get_child_with_siblings_for_test::(&internal_node_key, i.into(), None).unwrap(), ( Some(leaf1_node_key.clone()), vec![hash4.into(), hash_x4.into(), hash_x1.into()] @@ -391,7 +391,7 @@ proptest! { } prop_assert_eq!( - internal_node.get_child_with_siblings::(&internal_node_key, 2.into(), None).unwrap(), + internal_node.get_child_with_siblings_for_test::(&internal_node_key, 2.into(), None).unwrap(), ( Some(internal2_node_key), vec![ @@ -404,7 +404,7 @@ proptest! { ); prop_assert_eq!( - internal_node.get_child_with_siblings::(&internal_node_key, 3.into(), None).unwrap(), + internal_node.get_child_with_siblings_for_test::(&internal_node_key, 3.into(), None).unwrap(), ( None, @@ -414,7 +414,7 @@ proptest! { for i in 4..6 { prop_assert_eq!( - internal_node.get_child_with_siblings::(&internal_node_key, i.into(), None).unwrap(), + internal_node.get_child_with_siblings_for_test::(&internal_node_key, i.into(), None).unwrap(), ( None, vec![hash4.into(), hash_x2.into(), hash_x3.into()] @@ -423,7 +423,7 @@ proptest! { } prop_assert_eq!( - internal_node.get_child_with_siblings::(&internal_node_key, 6.into(), None).unwrap(), + internal_node.get_child_with_siblings_for_test::(&internal_node_key, 6.into(), None).unwrap(), ( None, vec![ @@ -436,7 +436,7 @@ proptest! { ); prop_assert_eq!( - internal_node.get_child_with_siblings::(&internal_node_key, 7.into(), None).unwrap(), + internal_node.get_child_with_siblings_for_test::(&internal_node_key, 7.into(), None).unwrap(), ( Some(internal3_node_key), vec![ @@ -450,7 +450,7 @@ proptest! { for i in 8..16 { prop_assert_eq!( - internal_node.get_child_with_siblings::(&internal_node_key, i.into(), None).unwrap(), + internal_node.get_child_with_siblings_for_test::(&internal_node_key, i.into(), None).unwrap(), (Some(leaf4_node_key.clone()), vec![hash_x5.into()]) ); } @@ -525,7 +525,7 @@ fn test_internal_hash_and_proof() { for i in 0..4 { assert_eq!( internal_node - .get_child_with_siblings::( + .get_child_with_siblings_for_test::( &internal_node_key, i.into(), None @@ -537,7 +537,11 @@ fn test_internal_hash_and_proof() { assert_eq!( internal_node - .get_child_with_siblings::(&internal_node_key, index1, None) + .get_child_with_siblings_for_test::( + &internal_node_key, + index1, + None + ) .unwrap(), (Some(child1_node_key), vec![ hash_x6.into(), @@ -549,7 +553,7 @@ fn test_internal_hash_and_proof() { assert_eq!( internal_node - .get_child_with_siblings::( + .get_child_with_siblings_for_test::( &internal_node_key, 5.into(), None @@ -565,7 +569,7 @@ fn test_internal_hash_and_proof() { for i in 6..8 { assert_eq!( internal_node - .get_child_with_siblings::( + .get_child_with_siblings_for_test::( &internal_node_key, i.into(), None @@ -582,7 +586,7 @@ fn test_internal_hash_and_proof() { for i in 8..12 { assert_eq!( internal_node - .get_child_with_siblings::( + .get_child_with_siblings_for_test::( &internal_node_key, i.into(), None @@ -595,7 +599,7 @@ fn test_internal_hash_and_proof() { for i in 12..14 { assert_eq!( internal_node - .get_child_with_siblings::( + .get_child_with_siblings_for_test::( &internal_node_key, i.into(), None @@ -610,7 +614,7 @@ fn test_internal_hash_and_proof() { } assert_eq!( internal_node - .get_child_with_siblings::( + .get_child_with_siblings_for_test::( &internal_node_key, 14.into(), None @@ -625,7 +629,11 @@ fn test_internal_hash_and_proof() { ); assert_eq!( internal_node - .get_child_with_siblings::(&internal_node_key, index2, None) + .get_child_with_siblings_for_test::( + &internal_node_key, + index2, + None + ) .unwrap(), (Some(child2_node_key), vec![ hash_x3.into(), @@ -701,7 +709,7 @@ fn test_internal_hash_and_proof() { assert_eq!( internal_node - .get_child_with_siblings::( + .get_child_with_siblings_for_test::( &internal_node_key, 0.into(), None @@ -717,7 +725,7 @@ fn test_internal_hash_and_proof() { assert_eq!( internal_node - .get_child_with_siblings::( + .get_child_with_siblings_for_test::( &internal_node_key, 1.into(), None @@ -734,7 +742,7 @@ fn test_internal_hash_and_proof() { for i in 2..4 { assert_eq!( internal_node - .get_child_with_siblings::( + .get_child_with_siblings_for_test::( &internal_node_key, i.into(), None @@ -751,7 +759,7 @@ fn test_internal_hash_and_proof() { for i in 4..6 { assert_eq!( internal_node - .get_child_with_siblings::( + .get_child_with_siblings_for_test::( &internal_node_key, i.into(), None @@ -767,7 +775,7 @@ fn test_internal_hash_and_proof() { assert_eq!( internal_node - .get_child_with_siblings::( + .get_child_with_siblings_for_test::( &internal_node_key, 6.into(), None @@ -783,7 +791,7 @@ fn test_internal_hash_and_proof() { assert_eq!( internal_node - .get_child_with_siblings::( + .get_child_with_siblings_for_test::( &internal_node_key, 7.into(), None @@ -800,7 +808,7 @@ fn test_internal_hash_and_proof() { for i in 8..16 { assert_eq!( internal_node - .get_child_with_siblings::( + .get_child_with_siblings_for_test::( &internal_node_key, i.into(), None @@ -984,7 +992,7 @@ proptest! { ) { for n in 0..16u8 { prop_assert_eq!( - node.get_child_with_siblings::(&node_key, n.into(), None).unwrap(), + node.get_child_with_siblings_for_test::(&node_key, n.into(), None).unwrap(), NaiveInternalNode::from_clever_node(&node).get_child_with_siblings(&node_key, n) ) }