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.
aptos-labs · May 15, 2024 · b265b3c · b265b3c
1 parent fb462d7
commit b265b3c
Show file tree

Hide file tree

Showing 15 changed files with 267 additions and 112 deletions.
diff --git 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<SparseMerkleProofExt> {
         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<StateValue>, 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<ExecutedTrees> {

diff --git a/storage/aptosdb/src/db/include/aptosdb_reader.rs b/storage/aptosdb/src/db/include/aptosdb_reader.rs
@@ -457,25 +457,27 @@ impl DbReader for AptosDB {
         &self,
         state_key: &StateKey,
         version: Version,
+        root_depth: usize,
     ) -> Result<SparseMerkleProofExt> {
         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)
         })
     }
 
     fn get_state_value_with_proof_by_version_ext(
         &self,
         state_store_key: &StateKey,
         version: Version,
+        root_depth: usize,
     ) -> Result<(Option<StateValue>, 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
@@ -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
@@ -165,10 +165,11 @@ impl DbReader for StateDb {
         &self,
         state_key: &StateKey,
         version: Version,
+        root_depth: usize,
     ) -> Result<SparseMerkleProofExt> {
         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<StateValue>, 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,19 +245,21 @@ impl DbReader for StateStore {
         &self,
         state_key: &StateKey,
         version: Version,
+        root_depth: usize,
     ) -> Result<SparseMerkleProofExt> {
         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
     fn get_state_value_with_proof_by_version_ext(
         &self,
         state_key: &StateKey,
         version: Version,
+        root_depth: usize,
     ) -> Result<(Option<StateValue>, 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
@@ -717,14 +717,15 @@ 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()))
     }
 
     pub fn get_with_proof_ext(
         &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
@@ -597,6 +597,8 @@ impl InternalNode {
         node_key: &NodeKey,
         n: Nibble,
         reader: Option<&R>,
+        root_depth: usize,
+        target_depth: usize,
     ) -> Result<(Option<NodeKey>, Vec<NodeInProof>)> {
         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<K: crate::Key, R: TreeReader<K>>(
+        &self,
+        node_key: &NodeKey,
+        n: Nibble,
+        reader: Option<&R>,
+    ) -> Result<(Option<NodeKey>, Vec<NodeInProof>)> {
+        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`