fix iteration over proofs

paradigmxyz · Sep 26, 2024 · ebc9944 · ebc9944
1 parent 45bae81
commit ebc9944
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 10 deletions.
diff --git a/crates/trie/common/src/proofs.rs b/crates/trie/common/src/proofs.rs
@@ -8,6 +8,7 @@ use alloy_trie::{
     proof::{verify_proof, ProofNodes, ProofVerificationError},
     EMPTY_ROOT_HASH,
 };
+use itertools::Itertools;
 use reth_primitives_traits::{constants::KECCAK_EMPTY, Account};
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
@@ -36,8 +37,8 @@ impl MultiProof {
         // Retrieve the account proof.
         let proof = self
             .account_subtree
-            .iter()
-            .filter(|(path, _)| nibbles.starts_with(path))
+            .matching_nodes_iter(&nibbles)
+            .sorted_by(|a, b| a.0.cmp(&b.0))
             .map(|(_, node)| node.clone())
             .collect::<Vec<_>>();
 
@@ -99,8 +100,8 @@ impl StorageMultiProof {
         // Retrieve the storage proof.
         let proof = self
             .subtree
-            .iter()
-            .filter(|(path, _)| nibbles.starts_with(path))
+            .matching_nodes_iter(&nibbles)
+            .sorted_by(|a, b| a.0.cmp(&b.0))
             .map(|(_, node)| node.clone())
             .collect::<Vec<_>>();
 

diff --git a/crates/trie/trie/src/witness.rs b/crates/trie/trie/src/witness.rs
@@ -120,23 +120,25 @@ where
                 None
             };
             let key = Nibbles::unpack(hashed_address);
-            let proof = account_multiproof.account_subtree.iter().filter(|e| key.starts_with(e.0));
-            account_trie_nodes.extend(self.target_nodes(key.clone(), value, proof)?);
+            account_trie_nodes.extend(self.target_nodes(
+                key.clone(),
+                value,
+                account_multiproof.account_subtree.matching_nodes_iter(&key),
+            )?);
 
             // Gather and record storage trie nodes for this account.
             let mut storage_trie_nodes = BTreeMap::default();
             let storage = state.storages.get(&hashed_address);
             for hashed_slot in hashed_slots {
-                let slot_key = Nibbles::unpack(hashed_slot);
+                let slot_nibbles = Nibbles::unpack(hashed_slot);
                 let slot_value = storage
                     .and_then(|s| s.storage.get(&hashed_slot))
                     .filter(|v| !v.is_zero())
                     .map(|v| alloy_rlp::encode_fixed_size(v).to_vec());
-                let proof = storage_multiproof.subtree.iter().filter(|e| slot_key.starts_with(e.0));
                 storage_trie_nodes.extend(self.target_nodes(
-                    slot_key.clone(),
+                    slot_nibbles.clone(),
                     slot_value,
-                    proof,
+                    storage_multiproof.subtree.matching_nodes_iter(&slot_nibbles),
                 )?);
             }