diff --git a/storage/jellyfish-merkle/src/lib.rs b/storage/jellyfish-merkle/src/lib.rs index f08a97305f86ba..344ef1edbae0ed 100644 --- 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 index ddb8916d224977..9552e61ed03f66 100644 --- 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, 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. @@ -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>( + &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 e75e05eac78fab..fde5ddbc8fdf46 100644 --- a/storage/jellyfish-merkle/src/node_type/node_type_test.rs +++ b/storage/jellyfish-merkle/src/node_type/node_type_test.rs @@ -199,13 +199,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()]) ); } @@ -246,14 +246,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![ @@ -267,7 +267,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![ @@ -281,7 +281,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()]) ); } @@ -320,21 +320,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()]) ); } @@ -385,7 +385,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()] @@ -394,7 +394,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![ @@ -407,7 +407,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, @@ -417,7 +417,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()] @@ -426,7 +426,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![ @@ -439,7 +439,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![ @@ -453,7 +453,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()]) ); } @@ -528,7 +528,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 @@ -540,7 +540,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(), @@ -552,7 +556,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 @@ -568,7 +572,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 @@ -585,7 +589,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 @@ -598,7 +602,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 @@ -613,7 +617,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 @@ -628,7 +632,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(), @@ -704,7 +712,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 @@ -720,7 +728,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 @@ -737,7 +745,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 @@ -754,7 +762,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 @@ -770,7 +778,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 @@ -786,7 +794,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 @@ -803,7 +811,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 @@ -987,7 +995,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) ) }