perf(trie): storage multiproof overallocation (paradigmxyz#12959)

mouseless0x · Dec 2, 2024 · 04f8c58 · 04f8c58
1 parent c2ab690
commit 04f8c58
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/crates/trie/parallel/src/proof.rs b/crates/trie/parallel/src/proof.rs
@@ -203,7 +203,11 @@ where
                     account.encode(&mut account_rlp as &mut dyn BufMut);
 
                     hash_builder.add_leaf(Nibbles::unpack(hashed_address), &account_rlp);
-                    storages.insert(hashed_address, storage_multiproof);
+
+                    // We might be adding leaves that are not necessarily our proof targets.
+                    if targets.contains_key(&hashed_address) {
+                        storages.insert(hashed_address, storage_multiproof);
+                    }
                 }
             }
         }

diff --git a/crates/trie/trie/src/proof.rs b/crates/trie/trie/src/proof.rs
@@ -115,19 +115,20 @@ where
                     hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
                 }
                 TrieElement::Leaf(hashed_address, account) => {
+                    let proof_targets = targets.remove(&hashed_address);
+                    let leaf_is_proof_target = proof_targets.is_some();
                     let storage_prefix_set = self
                         .prefix_sets
                         .storage_prefix_sets
                         .remove(&hashed_address)
                         .unwrap_or_default();
-                    let proof_targets = targets.remove(&hashed_address).unwrap_or_default();
                     let storage_multiproof = StorageProof::new_hashed(
                         self.trie_cursor_factory.clone(),
                         self.hashed_cursor_factory.clone(),
                         hashed_address,
                     )
                     .with_prefix_set_mut(storage_prefix_set)
-                    .storage_multiproof(proof_targets)?;
+                    .storage_multiproof(proof_targets.unwrap_or_default())?;
 
                     // Encode account
                     account_rlp.clear();
@@ -136,8 +137,11 @@ where
 
                     hash_builder.add_leaf(Nibbles::unpack(hashed_address), &account_rlp);
 
-                    // Overwrite storage multiproof.
-                    storages.insert(hashed_address, storage_multiproof);
+                    // We might be adding leaves that are not necessarily our proof targets.
+                    if leaf_is_proof_target {
+                        // Overwrite storage multiproof.
+                        storages.insert(hashed_address, storage_multiproof);
+                    }
                 }
             }
         }