From b265b3c0e7c9db13d1747c8ee06795c3642271f0 Mon Sep 17 00:00:00 2001 From: aldenhu Date: Tue, 14 May 2024 00:17:24 +0000 Subject: [PATCH] partial state proofs During execution, much of the proof is known info already on the in-mem SMT. This changes it so that only the known parts of the proofs are requested, saving much duplicated hashing required to calculate the internal (binary SMT) nodes compressed out by the JMT node. --- storage/aptosdb/src/db/fake_aptosdb.rs | 6 +- .../aptosdb/src/db/include/aptosdb_reader.rs | 6 +- storage/aptosdb/src/state_merkle_db.rs | 3 +- storage/aptosdb/src/state_store/mod.rs | 12 ++-- storage/jellyfish-merkle/src/lib.rs | 18 ++++- storage/jellyfish-merkle/src/node_type/mod.rs | 42 +++++++---- .../src/node_type/node_type_test.rs | 72 ++++++++++--------- storage/scratchpad/src/sparse_merkle/mod.rs | 29 ++++---- .../src/sparse_merkle/sparse_merkle_test.rs | 70 ++++++++++++++---- .../scratchpad/src/sparse_merkle/updater.rs | 13 ++-- .../src/async_proof_fetcher.rs | 22 +++--- .../src/cached_state_view.rs | 10 ++- storage/storage-interface/src/lib.rs | 4 +- storage/storage-interface/src/mock.rs | 1 + types/src/proof/definition.rs | 71 +++++++++++++++--- 15 files changed, 267 insertions(+), 112 deletions(-) diff --git a/storage/aptosdb/src/db/fake_aptosdb.rs b/storage/aptosdb/src/db/fake_aptosdb.rs index 918c615d13105a..ad008f28f8ebef 100644 --- a/storage/aptosdb/src/db/fake_aptosdb.rs +++ b/storage/aptosdb/src/db/fake_aptosdb.rs @@ -814,18 +814,20 @@ impl DbReader for FakeAptosDB { &self, state_key: &StateKey, version: Version, + root_depth: usize, ) -> Result { self.inner - .get_state_proof_by_version_ext(state_key, version) + .get_state_proof_by_version_ext(state_key, version, root_depth) } fn get_state_value_with_proof_by_version_ext( &self, state_key: &StateKey, version: Version, + root_depth: usize, ) -> Result<(Option, SparseMerkleProofExt)> { self.inner - .get_state_value_with_proof_by_version_ext(state_key, version) + .get_state_value_with_proof_by_version_ext(state_key, version, root_depth) } fn get_latest_executed_trees(&self) -> Result { diff --git a/storage/aptosdb/src/db/include/aptosdb_reader.rs b/storage/aptosdb/src/db/include/aptosdb_reader.rs index 3886bf05f256cf..49980be1f3e941 100644 --- a/storage/aptosdb/src/db/include/aptosdb_reader.rs +++ b/storage/aptosdb/src/db/include/aptosdb_reader.rs @@ -457,12 +457,13 @@ impl DbReader for AptosDB { &self, state_key: &StateKey, version: Version, + root_depth: usize, ) -> Result { gauged_api("get_state_proof_by_version_ext", || { self.error_if_state_merkle_pruned("State merkle", version)?; self.state_store - .get_state_proof_by_version_ext(state_key, version) + .get_state_proof_by_version_ext(state_key, version, root_depth) }) } @@ -470,12 +471,13 @@ impl DbReader for AptosDB { &self, state_store_key: &StateKey, version: Version, + root_depth: usize, ) -> Result<(Option, SparseMerkleProofExt)> { gauged_api("get_state_value_with_proof_by_version_ext", || { self.error_if_state_merkle_pruned("State merkle", version)?; self.state_store - .get_state_value_with_proof_by_version_ext(state_store_key, version) + .get_state_value_with_proof_by_version_ext(state_store_key, version, root_depth) }) } diff --git a/storage/aptosdb/src/state_merkle_db.rs b/storage/aptosdb/src/state_merkle_db.rs index 468d492af30547..d1a9b9333d98ab 100644 --- a/storage/aptosdb/src/state_merkle_db.rs +++ b/storage/aptosdb/src/state_merkle_db.rs @@ -249,12 +249,13 @@ impl StateMerkleDb { &self, state_key: &StateKey, version: Version, + root_depth: usize, ) -> Result<( Option<(HashValue, (StateKey, Version))>, SparseMerkleProofExt, )> { JellyfishMerkleTree::new(self) - .get_with_proof_ext(state_key.hash(), version) + .get_with_proof_ext(state_key.hash(), version, root_depth) .map_err(Into::into) } diff --git a/storage/aptosdb/src/state_store/mod.rs b/storage/aptosdb/src/state_store/mod.rs index e8e7eb01f09fb4..49111a2f3fd802 100644 --- a/storage/aptosdb/src/state_store/mod.rs +++ b/storage/aptosdb/src/state_store/mod.rs @@ -165,10 +165,11 @@ impl DbReader for StateDb { &self, state_key: &StateKey, version: Version, + root_depth: usize, ) -> Result { let (_, proof) = self .state_merkle_db - .get_with_proof_ext(state_key, version)?; + .get_with_proof_ext(state_key, version, root_depth)?; Ok(proof) } @@ -177,10 +178,11 @@ impl DbReader for StateDb { &self, state_key: &StateKey, version: Version, + root_depth: usize, ) -> Result<(Option, SparseMerkleProofExt)> { let (leaf_data, proof) = self .state_merkle_db - .get_with_proof_ext(state_key, version)?; + .get_with_proof_ext(state_key, version, root_depth)?; Ok(( match leaf_data { Some((_, (key, version))) => Some(self.expect_value_by_version(&key, version)?), @@ -243,9 +245,10 @@ impl DbReader for StateStore { &self, state_key: &StateKey, version: Version, + root_depth: usize, ) -> Result { self.deref() - .get_state_proof_by_version_ext(state_key, version) + .get_state_proof_by_version_ext(state_key, version, root_depth) } /// Get the state value with proof extension given the state key and version @@ -253,9 +256,10 @@ impl DbReader for StateStore { &self, state_key: &StateKey, version: Version, + root_depth: usize, ) -> Result<(Option, SparseMerkleProofExt)> { self.deref() - .get_state_value_with_proof_by_version_ext(state_key, version) + .get_state_value_with_proof_by_version_ext(state_key, version, root_depth) } } diff --git a/storage/jellyfish-merkle/src/lib.rs b/storage/jellyfish-merkle/src/lib.rs index 1d10cb7e0db4b2..415f492d03d2a8 100644 --- a/storage/jellyfish-merkle/src/lib.rs +++ b/storage/jellyfish-merkle/src/lib.rs @@ -717,7 +717,7 @@ where key: HashValue, version: Version, ) -> Result<(Option<(HashValue, (K, Version))>, SparseMerkleProof)> { - self.get_with_proof_ext(key, version) + self.get_with_proof_ext(key, version, 0) .map(|(value, proof_ext)| (value, proof_ext.into())) } @@ -725,6 +725,7 @@ where &self, key: HashValue, version: Version, + target_root_depth: usize, ) -> 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); @@ -764,23 +765,34 @@ where &next_node_key, queried_child_index, Some(self.reader), + nibble_depth * 4, + target_root_depth, )?; siblings.append(&mut siblings_in_internal); next_node_key = match child_node_key { Some(node_key) => node_key, None => { - return Ok((None, SparseMerkleProofExt::new(None, siblings))); + let bottom_depth = target_root_depth + siblings.len(); + return Ok(( + None, + SparseMerkleProofExt::new_partial(None, siblings, bottom_depth), + )); }, }; }, Node::Leaf(leaf_node) => { + let bottom_depth = target_root_depth + siblings.len(); return Ok(( if leaf_node.account_key() == key { Some((leaf_node.value_hash(), leaf_node.value_index().clone())) } else { None }, - SparseMerkleProofExt::new(Some(leaf_node.into()), siblings), + SparseMerkleProofExt::new_partial( + Some(leaf_node.into()), + siblings, + bottom_depth, + ), )); }, Node::Null => { diff --git a/storage/jellyfish-merkle/src/node_type/mod.rs b/storage/jellyfish-merkle/src/node_type/mod.rs index 26b1740fb95102..dd6ae2352c51e4 100644 --- a/storage/jellyfish-merkle/src/node_type/mod.rs +++ b/storage/jellyfish-merkle/src/node_type/mod.rs @@ -597,6 +597,8 @@ impl InternalNode { node_key: &NodeKey, n: Nibble, reader: Option<&R>, + root_depth: usize, + target_depth: usize, ) -> Result<(Option, Vec)> { assert!(self.leaf_count > 1); @@ -610,20 +612,26 @@ impl InternalNode { let width = 1 << h; 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( - sibling_half_start, - width, - (existence_bitmap, leaf_bitmap), - (reader, node_key), - )?); - } else { - siblings.push( - self.merkle_hash(sibling_half_start, width, (existence_bitmap, leaf_bitmap)) + let depth = root_depth + 3 - h as usize; + if depth >= target_depth { + if let Some(reader) = reader { + siblings.push(self.gen_node_in_proof( + sibling_half_start, + width, + (existence_bitmap, leaf_bitmap), + (reader, node_key), + )?); + } else { + siblings.push( + self.merkle_hash( + sibling_half_start, + width, + (existence_bitmap, leaf_bitmap), + ) .into(), - ); + ); + } } - let (range_existence_bitmap, range_leaf_bitmap) = Self::range_bitmaps(child_half_start, width, (existence_bitmap, leaf_bitmap)); @@ -660,6 +668,16 @@ impl InternalNode { } unreachable!("Impossible to get here without returning even at the lowest level.") } + + #[cfg(test)] + pub fn get_child_with_siblings_for_test>( + &self, + node_key: &NodeKey, + n: Nibble, + reader: Option<&R>, + ) -> Result<(Option, Vec)> { + self.get_child_with_siblings(node_key, n, reader, 0, 0) + } } /// 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 ab1f7e1f21439a..77cab6a7f5df0b 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) ) } diff --git a/storage/scratchpad/src/sparse_merkle/mod.rs b/storage/scratchpad/src/sparse_merkle/mod.rs index f3a0c832e4d52c..f4ec62770bf003 100644 --- a/storage/scratchpad/src/sparse_merkle/mod.rs +++ b/storage/scratchpad/src/sparse_merkle/mod.rs @@ -441,20 +441,13 @@ pub enum StateStoreStatus { /// The entry exists in the tree, therefore we can give its value. ExistsInScratchPad(V), - /// The entry does not exist in the tree, but exists in DB. This happens when the search - /// reaches a leaf node that has the requested account, but the node has only the value hash - /// because it was loaded into memory as part of a non-inclusion proof. When we go to DB we - /// don't need to traverse the tree to find the same leaf, instead we can use the value hash to - /// look up the entry content directly. - ExistsInDB, - /// The entry does not exist in either the tree or DB. This happens when the search reaches /// an empty node, or a leaf node that has a different account. DoesNotExist, - /// We do not know if this entry exists or not and need to go to DB to find out. This happens - /// when the search reaches a subtree node. - Unknown, + /// Tree nodes only exist until `depth` on the route from the root to the leaf address, needs + /// to check the DB for the rest. + UnknownSubtreeRoot { hash: HashValue, depth: usize }, } /// In the entire lifetime of this, in-mem nodes won't be dropped because a reference to the oldest @@ -530,13 +523,20 @@ where pub fn get(&self, key: HashValue) -> StateStoreStatus { let mut subtree = self.smt.root_weak(); let mut bits = key.iter_bits(); + let mut next_depth = 0; loop { + next_depth += 1; match subtree { SubTree::Empty => return StateStoreStatus::DoesNotExist, - SubTree::NonEmpty { .. } => { + SubTree::NonEmpty { hash, root: _ } => { match subtree.get_node_if_in_mem(self.base_generation) { - None => return StateStoreStatus::Unknown, + None => { + return StateStoreStatus::UnknownSubtreeRoot { + hash, + depth: next_depth - 1, + } + }, Some(node) => match node.inner() { NodeInner::Internal(internal_node) => { subtree = if bits.next().expect("Tree is too deep.") { @@ -552,7 +552,10 @@ where Some(value) => StateStoreStatus::ExistsInScratchPad( value.as_ref().clone(), ), - None => StateStoreStatus::ExistsInDB, + None => StateStoreStatus::UnknownSubtreeRoot { + hash, + depth: next_depth - 1, + }, } } else { StateStoreStatus::DoesNotExist diff --git a/storage/scratchpad/src/sparse_merkle/sparse_merkle_test.rs b/storage/scratchpad/src/sparse_merkle/sparse_merkle_test.rs index bec1430c6c0f01..dc3c2780e165ec 100644 --- a/storage/scratchpad/src/sparse_merkle/sparse_merkle_test.rs +++ b/storage/scratchpad/src/sparse_merkle/sparse_merkle_test.rs @@ -247,7 +247,10 @@ fn test_update_256_siblings_in_proof() { new_smt.get(key1), StateStoreStatus::ExistsInScratchPad(new_value1) ); - assert_eq!(new_smt.get(key2), StateStoreStatus::ExistsInDB); + assert_eq!(new_smt.get(key2), StateStoreStatus::UnknownSubtreeRoot { + hash: leaf2.hash(), + depth: 256 + }); } #[test] @@ -320,9 +323,18 @@ fn test_update() { // y key3 (unknown) // / \ // x key4 - assert_eq!(smt1.get(key1), StateStoreStatus::Unknown); - assert_eq!(smt1.get(key2), StateStoreStatus::Unknown); - assert_eq!(smt1.get(key3), StateStoreStatus::ExistsInDB); + assert_eq!(smt1.get(key1), StateStoreStatus::UnknownSubtreeRoot { + hash: x_hash, + depth: 2 + }); + assert_eq!(smt1.get(key2), StateStoreStatus::UnknownSubtreeRoot { + hash: x_hash, + depth: 2 + }); + assert_eq!(smt1.get(key3), StateStoreStatus::UnknownSubtreeRoot { + hash: leaf3.hash(), + depth: 1 + }); assert_eq!( smt1.get(key4), StateStoreStatus::ExistsInScratchPad(value4.clone()) @@ -359,8 +371,14 @@ fn test_update() { // / \ // key2(indb) key4 (weak data) assert_eq!(smt2.get(key1), StateStoreStatus::DoesNotExist,); - assert_eq!(smt2.get(key2), StateStoreStatus::ExistsInDB); - assert_eq!(smt2.get(key3), StateStoreStatus::ExistsInDB); + assert_eq!(smt2.get(key2), StateStoreStatus::UnknownSubtreeRoot { + hash: leaf2.hash(), + depth: 2 + }); + assert_eq!(smt2.get(key3), StateStoreStatus::UnknownSubtreeRoot { + hash: leaf3.hash(), + depth: 1 + }); assert_eq!(smt2.get(key4), StateStoreStatus::ExistsInScratchPad(value4)); // Verify root hash. @@ -383,8 +401,14 @@ fn test_update() { // y' key3 (indb, weak) // / \ // (weak) x key4 - assert_eq!(smt22.get(key2), StateStoreStatus::Unknown); - assert_eq!(smt22.get(key3), StateStoreStatus::ExistsInDB); + assert_eq!(smt22.get(key2), StateStoreStatus::UnknownSubtreeRoot { + hash: x_hash, + depth: 2 + }); + assert_eq!(smt22.get(key3), StateStoreStatus::UnknownSubtreeRoot { + hash: leaf3.hash(), + depth: 1 + }); assert_eq!( smt22.get(key4), StateStoreStatus::ExistsInScratchPad(value4.clone()) @@ -396,15 +420,33 @@ fn test_update() { // For smt2, no key should be available since smt2 was constructed by deleting key1. assert_eq!(smt2.get(key1), StateStoreStatus::DoesNotExist); - assert_eq!(smt2.get(key2), StateStoreStatus::ExistsInDB); - assert_eq!(smt2.get(key3), StateStoreStatus::ExistsInDB); - assert_eq!(smt2.get(key4), StateStoreStatus::ExistsInDB); + assert_eq!(smt2.get(key2), StateStoreStatus::UnknownSubtreeRoot { + hash: leaf2.hash(), + depth: 2 + }); + assert_eq!(smt2.get(key3), StateStoreStatus::UnknownSubtreeRoot { + hash: leaf3.hash(), + depth: 1 + }); + assert_eq!(smt2.get(key4), StateStoreStatus::UnknownSubtreeRoot { + hash: leaf4_hash, + depth: 2 + }); // For smt22, only key4 should be available since smt22 was constructed by updating smt1 with // key4. - assert_eq!(smt22.get(key1), StateStoreStatus::Unknown); - assert_eq!(smt22.get(key2), StateStoreStatus::Unknown); - assert_eq!(smt22.get(key3), StateStoreStatus::ExistsInDB); + assert_eq!(smt22.get(key1), StateStoreStatus::UnknownSubtreeRoot { + hash: x_hash, + depth: 2 + }); + assert_eq!(smt22.get(key2), StateStoreStatus::UnknownSubtreeRoot { + hash: x_hash, + depth: 2 + }); + assert_eq!(smt22.get(key3), StateStoreStatus::UnknownSubtreeRoot { + hash: leaf3.hash(), + depth: 1 + }); assert_eq!( smt22.get(key4), StateStoreStatus::ExistsInScratchPad(value4) diff --git a/storage/scratchpad/src/sparse_merkle/updater.rs b/storage/scratchpad/src/sparse_merkle/updater.rs index f923687680eaac..95dd16a3247c93 100644 --- a/storage/scratchpad/src/sparse_merkle/updater.rs +++ b/storage/scratchpad/src/sparse_merkle/updater.rs @@ -21,7 +21,7 @@ use std::cmp::Ordering; type Result = std::result::Result; type InMemSubTree = super::node::SubTree; -type InMemInternal = super::node::InternalNode; +type InMemInternal = InternalNode; #[derive(Clone)] enum InMemSubTreeInfo { @@ -133,7 +133,7 @@ impl<'a, V: Clone + CryptoHash + Send + Sync + 'static> SubTreeInfo<'a, V> { } fn new_on_proof_path(proof: &'a SparseMerkleProofExt, depth: usize) -> Self { - match proof.siblings().len().cmp(&depth) { + match proof.bottom_depth().cmp(&depth) { Ordering::Greater => Self::Persisted(PersistedSubTreeInfo::ProofPathInternal { proof }), Ordering::Equal => match proof.leaf() { Some(leaf) => Self::new_proof_leaf(leaf), @@ -151,10 +151,10 @@ impl<'a, V: Clone + CryptoHash + Send + Sync + 'static> SubTreeInfo<'a, V> { let proof = proof_reader .get_proof(a_descendant_key) .ok_or(UpdateError::MissingProof)?; - if depth > proof.siblings().len() { + if depth > proof.bottom_depth() { return Err(UpdateError::ShortProof { key: a_descendant_key, - num_siblings: proof.siblings().len(), + num_siblings: proof.bottom_depth(), depth, }); } @@ -244,9 +244,8 @@ impl<'a, V: Clone + CryptoHash + Send + Sync + 'static> SubTreeInfo<'a, V> { swap_if(myself, SubTreeInfo::new_empty(), key.bit(depth)) }, PersistedSubTreeInfo::ProofPathInternal { proof } => { - let siblings = proof.siblings(); - assert!(siblings.len() > depth); - let sibling_child = SubTreeInfo::new_proof_sibling(&siblings[depth]); + let sibling_child = + SubTreeInfo::new_proof_sibling(proof.sibling_at_depth(depth + 1).unwrap()); let on_path_child = SubTreeInfo::new_on_proof_path(proof, depth + 1); swap_if(on_path_child, sibling_child, a_descendent_key.bit(depth)) }, diff --git a/storage/storage-interface/src/async_proof_fetcher.rs b/storage/storage-interface/src/async_proof_fetcher.rs index fbc344eb8f9f53..87fb1d0dcc0861 100644 --- a/storage/storage-interface/src/async_proof_fetcher.rs +++ b/storage/storage-interface/src/async_proof_fetcher.rs @@ -59,7 +59,8 @@ impl AsyncProofFetcher { &self, state_key: &StateKey, version: Version, - root_hash: Option, + subtree_root_depth: usize, + subtree_root_hash: Option, ) -> Result> { let _timer = TIMER .with_label_values(&["async_proof_fetcher_fetch"]) @@ -70,7 +71,8 @@ impl AsyncProofFetcher { self.schedule_proof_read( state_key.clone(), version, - root_hash, + subtree_root_depth, + subtree_root_hash, version_and_value_opt.as_ref().map(|v| { let state_value = &v.1; state_value.hash() @@ -110,7 +112,8 @@ impl AsyncProofFetcher { &self, state_key: StateKey, version: Version, - root_hash: Option, + subtree_root_depth: usize, + subtree_root_hash: Option, value_hash: Option, ) { let _timer = TIMER @@ -121,19 +124,20 @@ impl AsyncProofFetcher { let data_sender = self.data_sender.clone(); IO_POOL.execute(move || { let proof = reader - .get_state_proof_by_version_ext(&state_key, version) + .get_state_proof_by_version_ext(&state_key, version, subtree_root_depth) .expect("Proof reading should succeed."); // NOTE: Drop the reader here to make sure reader has shorter lifetime than the async // proof fetcher. drop(reader); - if let Some(root_hash) = root_hash { + if let Some(subtree_root_hash) = subtree_root_hash { proof - .verify_by_hash(root_hash, state_key.hash(), value_hash) + .verify_by_hash(subtree_root_hash, state_key.hash(), value_hash) .map_err(|err| { anyhow!( - "Proof is invalid for key {:?} with state root hash {:?}, at version {}: {}.", + "Proof is invalid for key {:?} with subtree root hash {:?}, depth {}, at version {}: {}.", state_key, - root_hash, + subtree_root_hash, + subtree_root_depth, version, err ) @@ -171,7 +175,7 @@ mod tests { let state_key: StateKey = StateKey::raw(format!("test_key_{}", i).as_bytes()); expected_key_hashes.push(state_key.hash()); let result = fetcher - .fetch_state_value_with_version_and_schedule_proof_read(&state_key, 0, None) + .fetch_state_value_with_version_and_schedule_proof_read(&state_key, 0, 0, None) .expect("Should not fail."); let expected_value = StateValue::from(match state_key.inner() { StateKeyInner::Raw(key) => key.to_owned(), diff --git a/storage/storage-interface/src/cached_state_view.rs b/storage/storage-interface/src/cached_state_view.rs index 123e49d1b8bd56..9f52e6f1887b5e 100644 --- a/storage/storage-interface/src/cached_state_view.rs +++ b/storage/storage-interface/src/cached_state_view.rs @@ -225,14 +225,18 @@ impl CachedStateView { StateStoreStatus::DoesNotExist => (None, None), // No matter it is in db or unknown, we have to query from db since even the // former case, we don't have the blob data but only its hash. - StateStoreStatus::ExistsInDB | StateStoreStatus::Unknown => match self.snapshot { - Some((version, root_hash)) => { + StateStoreStatus::UnknownSubtreeRoot { + hash: subtree_root_hash, + depth: subtree_root_depth, + } => match self.snapshot { + Some((version, _root_hash)) => { let version_and_value_opt = self .proof_fetcher .fetch_state_value_with_version_and_schedule_proof_read( state_key, version, - Some(root_hash), + subtree_root_depth, + Some(subtree_root_hash), )?; match version_and_value_opt { Some((version, value)) => (Some(version), Some(value)), diff --git a/storage/storage-interface/src/lib.rs b/storage/storage-interface/src/lib.rs index ae8e43256106d4..8281c0522c3dac 100644 --- a/storage/storage-interface/src/lib.rs +++ b/storage/storage-interface/src/lib.rs @@ -360,6 +360,7 @@ pub trait DbReader: Send + Sync { &self, state_key: &StateKey, version: Version, + root_depth: usize, ) -> Result; /// Gets a state value by state key along with the proof, out of the ledger state indicated by the state @@ -374,6 +375,7 @@ pub trait DbReader: Send + Sync { &self, state_key: &StateKey, version: Version, + root_depth: usize, ) -> Result<(Option, SparseMerkleProofExt)>; /// Gets the latest ExecutedTrees no matter if db has been bootstrapped. @@ -474,7 +476,7 @@ pub trait DbReader: Send + Sync { state_key: &StateKey, version: Version, ) -> Result<(Option, SparseMerkleProof)> { - self.get_state_value_with_proof_by_version_ext(state_key, version) + self.get_state_value_with_proof_by_version_ext(state_key, version, 0) .map(|(value, proof_ext)| (value, proof_ext.into())) } } diff --git a/storage/storage-interface/src/mock.rs b/storage/storage-interface/src/mock.rs index 664e589150100b..5d3d367cc35b4e 100644 --- a/storage/storage-interface/src/mock.rs +++ b/storage/storage-interface/src/mock.rs @@ -27,6 +27,7 @@ impl DbReader for MockDbReaderWriter { &self, _state_key: &StateKey, _version: Version, + _root_depth: usize, ) -> Result { Ok(SparseMerkleProofExt::new(None, vec![])) } diff --git a/types/src/proof/definition.rs b/types/src/proof/definition.rs index 1d88b56c52a66e..2aeaffe3d13fdb 100644 --- a/types/src/proof/definition.rs +++ b/types/src/proof/definition.rs @@ -186,12 +186,32 @@ pub struct SparseMerkleProofExt { /// All siblings in this proof, including the default ones. Siblings are ordered from the root /// level to the bottom level. siblings: Vec, + /// depth of the leaf (or the None-leaf that proves there's no such leaf) + bottom_depth: usize, } impl SparseMerkleProofExt { /// Constructs a new `SparseMerkleProofExt` using leaf and a list of sibling nodes. pub fn new(leaf: Option, siblings: Vec) -> Self { - Self { leaf, siblings } + let bottom_depth = siblings.len(); + Self { + leaf, + siblings, + bottom_depth, + } + } + + pub fn new_partial( + leaf: Option, + siblings: Vec, + bottom_depth: usize, + ) -> Self { + assert!(bottom_depth >= siblings.len()); + Self { + leaf, + siblings, + bottom_depth, + } } /// Returns the leaf node in this proof. @@ -199,9 +219,26 @@ impl SparseMerkleProofExt { self.leaf } - /// Returns the list of siblings in this proof. - pub fn siblings(&self) -> &[NodeInProof] { - &self.siblings + pub fn sibling_at_depth(&self, depth: usize) -> Result<&NodeInProof> { + ensure!( + depth > self.root_depth() && depth <= self.bottom_depth, + "Proof between depth {} and {} does not cover depth {}", + self.root_depth(), + self.bottom_depth(), + depth, + ); + Ok(&self.siblings[depth - self.root_depth() - 1]) + } + + /// Returns the depth of the leaf (or the None-leaf that proves there's no such leaf). + pub fn bottom_depth(&self) -> usize { + self.bottom_depth + } + + /// Assuming a possible partial proof, returns the depth of the root of the subtree this + /// proof proves, i.e. this holds: `self.root_depth() + self.siblings.len() == self.depth` + pub fn root_depth(&self) -> usize { + self.bottom_depth - self.siblings.len() } pub fn verify( @@ -210,7 +247,11 @@ impl SparseMerkleProofExt { element_key: HashValue, element_value: Option<&V>, ) -> Result<()> { - SparseMerkleProof::from(self.clone()).verify(expected_root_hash, element_key, element_value) + self.verify_by_hash( + expected_root_hash, + element_key, + element_value.map(|v| v.hash()), + ) } pub fn verify_by_hash( @@ -219,10 +260,11 @@ impl SparseMerkleProofExt { element_key: HashValue, element_hash: Option, ) -> Result<()> { - SparseMerkleProof::from(self.clone()).verify_by_hash( + SparseMerkleProof::from(self.clone()).verify_by_hash_partial( expected_root_hash, element_key, element_hash, + self.root_depth(), ) } } @@ -278,11 +320,22 @@ impl SparseMerkleProof { expected_root_hash: HashValue, element_key: HashValue, element_hash: Option, + ) -> Result<()> { + self.verify_by_hash_partial(expected_root_hash, element_key, element_hash, 0) + } + + pub fn verify_by_hash_partial( + &self, + expected_root_hash: HashValue, + element_key: HashValue, + element_hash: Option, + root_depth: usize, ) -> Result<()> { ensure!( - self.siblings.len() <= HashValue::LENGTH_IN_BITS, - "Sparse Merkle Tree proof has more than {} ({}) siblings.", + self.siblings.len() + root_depth <= HashValue::LENGTH_IN_BITS, + "Sparse Merkle Tree proof has more than {} ({} + {}) siblings.", HashValue::LENGTH_IN_BITS, + root_depth, self.siblings.len(), ); @@ -354,7 +407,7 @@ impl SparseMerkleProof { element_key .iter_bits() .rev() - .skip(HashValue::LENGTH_IN_BITS - self.siblings.len()), + .skip(HashValue::LENGTH_IN_BITS - self.siblings.len() - root_depth), ) .fold(current_hash, |hash, (sibling_hash, bit)| { if bit {