feat(trie): expose multiproof via StateProofProvider (#10915)

crates/chain-state/src/in_memory.rs

@@ -839,7 +839,8 @@ mod tests {
         AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateRootProvider,
         StorageRootProvider,
     };
-    use reth_trie::{AccountProof, HashedStorage, TrieInput};
+    use reth_trie::{AccountProof, HashedStorage, MultiProof, TrieInput};
+    use std::collections::HashSet;
 
     fn create_mock_state(
         test_block_builder: &mut TestBlockBuilder,
@@ -952,6 +953,14 @@ mod tests {
             Ok(AccountProof::new(Address::random()))
         }
 
+        fn multiproof(
+            &self,
+            _input: TrieInput,
+            _targets: HashMap<B256, HashSet<B256>>,
+        ) -> ProviderResult<MultiProof> {
+            Ok(MultiProof::default())
+        }
+
         fn witness(
             &self,
             _input: TrieInput,

crates/chain-state/src/memory_overlay.rs

@@ -7,8 +7,13 @@ use reth_storage_api::{
     AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateProviderBox,
     StateRootProvider, StorageRootProvider,
 };
-use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, TrieInput};
-use std::{collections::HashMap, sync::OnceLock};
+use reth_trie::{
+    updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof, TrieInput,
+};
+use std::{
+    collections::{HashMap, HashSet},
+    sync::OnceLock,
+};
 
 /// A state provider that stores references to in-memory blocks along with their state as well as
 /// the historical state provider for fallback lookups.
@@ -147,6 +152,16 @@ impl StateProofProvider for MemoryOverlayStateProvider {
         self.historical.proof(input, address, slots)
     }
 
+    fn multiproof(
+        &self,
+        mut input: TrieInput,
+        targets: HashMap<B256, HashSet<B256>>,
+    ) -> ProviderResult<MultiProof> {
+        let MemoryOverlayTrieState { nodes, state } = self.trie_state().clone();
+        input.prepend_cached(nodes, state);
+        self.historical.multiproof(input, targets)
+    }
+
     fn witness(
         &self,
         mut input: TrieInput,

crates/revm/src/test_utils.rs

@@ -1,3 +1,5 @@
+use std::collections::HashSet;
+
 use crate::precompile::HashMap;
 use alloc::vec::Vec;
 use reth_primitives::{
@@ -8,7 +10,9 @@ use reth_storage_api::{
     StorageRootProvider,
 };
 use reth_storage_errors::provider::ProviderResult;
-use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, TrieInput};
+use reth_trie::{
+    updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof, TrieInput,
+};
 
 /// Mock state for testing
 #[derive(Debug, Default, Clone, Eq, PartialEq)]
@@ -110,6 +114,14 @@ impl StateProofProvider for StateProviderTest {
         unimplemented!("proof generation is not supported")
     }
 
+    fn multiproof(
+        &self,
+        _input: TrieInput,
+        _targets: HashMap<B256, HashSet<B256>>,
+    ) -> ProviderResult<MultiProof> {
+        unimplemented!("proof generation is not supported")
+    }
+
     fn witness(
         &self,
         _input: TrieInput,

crates/rpc/rpc-eth-types/src/cache/db.rs

@@ -67,6 +67,14 @@ impl<'a> reth_storage_api::StateProofProvider for StateProviderTraitObjWrapper<'
         self.0.proof(input, address, slots)
     }
 
+    fn multiproof(
+        &self,
+        input: reth_trie::TrieInput,
+        targets: std::collections::HashMap<B256, std::collections::HashSet<B256>>,
+    ) -> ProviderResult<reth_trie::MultiProof> {
+        self.0.multiproof(input, targets)
+    }
+
     fn witness(
         &self,
         input: reth_trie::TrieInput,

crates/storage/provider/src/providers/bundle_state_provider.rs

@@ -4,8 +4,10 @@ use crate::{
 use reth_primitives::{Account, Address, BlockNumber, Bytecode, Bytes, B256};
 use reth_storage_api::{StateProofProvider, StorageRootProvider};
 use reth_storage_errors::provider::ProviderResult;
-use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, TrieInput};
-use std::collections::HashMap;
+use reth_trie::{
+    updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof, TrieInput,
+};
+use std::collections::{HashMap, HashSet};
 
 /// A state provider that resolves to data from either a wrapped [`crate::ExecutionOutcome`]
 /// or an underlying state provider.
@@ -133,6 +135,16 @@ impl<SP: StateProvider, EDP: ExecutionDataProvider> StateProofProvider
         self.state_provider.proof(input, address, slots)
     }
 
+    fn multiproof(
+        &self,
+        mut input: reth_trie::TrieInput,
+        targets: HashMap<B256, HashSet<B256>>,
+    ) -> ProviderResult<MultiProof> {
+        let bundle_state = self.block_execution_data_provider.execution_outcome().state();
+        input.prepend(HashedPostState::from_bundle_state(&bundle_state.state));
+        self.state_provider.multiproof(input, targets)
+    }
+
     fn witness(
         &self,
         mut input: TrieInput,

crates/storage/provider/src/providers/state/historical.rs

@@ -17,13 +17,16 @@ use reth_storage_api::{StateProofProvider, StorageRootProvider};
 use reth_storage_errors::provider::ProviderResult;
 use reth_trie::{
     proof::Proof, updates::TrieUpdates, witness::TrieWitness, AccountProof, HashedPostState,
-    HashedStorage, StateRoot, StorageRoot, TrieInput,
+    HashedStorage, MultiProof, StateRoot, StorageRoot, TrieInput,
 };
 use reth_trie_db::{
     DatabaseHashedPostState, DatabaseHashedStorage, DatabaseProof, DatabaseStateRoot,
     DatabaseStorageRoot, DatabaseTrieWitness,
 };
-use std::{collections::HashMap, fmt::Debug};
+use std::{
+    collections::{HashMap, HashSet},
+    fmt::Debug,
+};
 
 /// State provider for a given block number which takes a tx reference.
 ///
@@ -344,6 +347,15 @@ impl<'b, TX: DbTx> StateProofProvider for HistoricalStateProviderRef<'b, TX> {
             .map_err(Into::<ProviderError>::into)
     }
 
+    fn multiproof(
+        &self,
+        mut input: TrieInput,
+        targets: HashMap<B256, HashSet<B256>>,
+    ) -> ProviderResult<MultiProof> {
+        input.prepend(self.revert_state()?);
+        Proof::overlay_multiproof(self.tx, input, targets).map_err(Into::<ProviderError>::into)
+    }
+
     fn witness(
         &self,
         mut input: TrieInput,

crates/storage/provider/src/providers/state/latest.rs

@@ -15,10 +15,10 @@ use reth_storage_api::{StateProofProvider, StorageRootProvider};
 use reth_storage_errors::provider::{ProviderError, ProviderResult};
 use reth_trie::{
     proof::Proof, updates::TrieUpdates, witness::TrieWitness, AccountProof, HashedPostState,
-    HashedStorage, StateRoot, StorageRoot, TrieInput,
+    HashedStorage, MultiProof, StateRoot, StorageRoot, TrieInput,
 };
 use reth_trie_db::{DatabaseProof, DatabaseStateRoot, DatabaseStorageRoot, DatabaseTrieWitness};
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 /// State provider over latest state that takes tx reference.
 #[derive(Debug)]
@@ -129,6 +129,14 @@ impl<'b, TX: DbTx> StateProofProvider for LatestStateProviderRef<'b, TX> {
             .map_err(Into::<ProviderError>::into)
     }
 
+    fn multiproof(
+        &self,
+        input: TrieInput,
+        targets: HashMap<B256, HashSet<B256>>,
+    ) -> ProviderResult<MultiProof> {
+        Proof::overlay_multiproof(self.tx, input, targets).map_err(Into::<ProviderError>::into)
+    }
+
     fn witness(
         &self,
         input: TrieInput,

crates/storage/provider/src/providers/state/macros.rs

@@ -52,6 +52,7 @@ macro_rules! delegate_provider_impls {
             }
             StateProofProvider $(where [$($generics)*])? {
                 fn proof(&self, input: reth_trie::TrieInput, address: reth_primitives::Address, slots: &[reth_primitives::B256]) -> reth_storage_errors::provider::ProviderResult<reth_trie::AccountProof>;
+                fn multiproof(&self, input: reth_trie::TrieInput, targets: std::collections::HashMap<reth_primitives::B256, std::collections::HashSet<reth_primitives::B256>>) -> reth_storage_errors::provider::ProviderResult<reth_trie::MultiProof>;
                 fn witness(&self, input: reth_trie::TrieInput, target: reth_trie::HashedPostState) -> reth_storage_errors::provider::ProviderResult<std::collections::HashMap<reth_primitives::B256, reth_primitives::Bytes>>;
             }
         );

crates/storage/provider/src/test_utils/mock.rs

@@ -24,10 +24,12 @@ use reth_storage_api::{
     DatabaseProviderFactory, StageCheckpointReader, StateProofProvider, StorageRootProvider,
 };
 use reth_storage_errors::provider::{ConsistentViewError, ProviderError, ProviderResult};
-use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, TrieInput};
+use reth_trie::{
+    updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof, TrieInput,
+};
 use revm::primitives::{BlockEnv, CfgEnvWithHandlerCfg};
 use std::{
-    collections::{BTreeMap, HashMap},
+    collections::{BTreeMap, HashMap, HashSet},
     ops::{RangeBounds, RangeInclusive},
     sync::Arc,
 };
@@ -624,6 +626,14 @@ impl StateProofProvider for MockEthProvider {
         Ok(AccountProof::new(address))
     }
 
+    fn multiproof(
+        &self,
+        _input: TrieInput,
+        _targets: HashMap<B256, HashSet<B256>>,
+    ) -> ProviderResult<MultiProof> {
+        Ok(MultiProof::default())
+    }
+
     fn witness(
         &self,
         _input: TrieInput,

crates/storage/provider/src/test_utils/noop.rs

@@ -1,5 +1,5 @@
 use std::{
-    collections::HashMap,
+    collections::{HashMap, HashSet},
     ops::{RangeBounds, RangeInclusive},
     sync::Arc,
 };
@@ -22,7 +22,9 @@ use reth_prune_types::{PruneCheckpoint, PruneSegment};
 use reth_stages_types::{StageCheckpoint, StageId};
 use reth_storage_api::{StateProofProvider, StorageRootProvider};
 use reth_storage_errors::provider::ProviderResult;
-use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, TrieInput};
+use reth_trie::{
+    updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof, TrieInput,
+};
 use revm::primitives::{BlockEnv, CfgEnvWithHandlerCfg};
 use tokio::sync::{broadcast, watch};
 
@@ -362,6 +364,14 @@ impl StateProofProvider for NoopProvider {
         Ok(AccountProof::new(address))
     }
 
+    fn multiproof(
+        &self,
+        _input: TrieInput,
+        _targets: HashMap<B256, HashSet<B256>>,
+    ) -> ProviderResult<MultiProof> {
+        Ok(MultiProof::default())
+    }
+
     fn witness(
         &self,
         _input: TrieInput,

crates/storage/storage-api/src/trie.rs

@@ -1,7 +1,9 @@
 use alloy_primitives::{Address, Bytes, B256};
 use reth_storage_errors::provider::ProviderResult;
-use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, TrieInput};
-use std::collections::HashMap;
+use reth_trie::{
+    updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof, TrieInput,
+};
+use std::collections::{HashMap, HashSet};
 
 /// A type that can compute the state root of a given post state.
 #[auto_impl::auto_impl(&, Box, Arc)]
@@ -56,6 +58,14 @@ pub trait StateProofProvider: Send + Sync {
         slots: &[B256],
     ) -> ProviderResult<AccountProof>;
 
+    /// Generate [`MultiProof`] for target hashed account and corresponding
+    /// hashed storage slot keys.
+    fn multiproof(
+        &self,
+        input: TrieInput,
+        targets: HashMap<B256, HashSet<B256>>,
+    ) -> ProviderResult<MultiProof>;
+
     /// Get trie witness for provided state.
     fn witness(
         &self,

crates/trie/common/src/proofs.rs

@@ -13,6 +13,8 @@ use serde::{Deserialize, Serialize};
 use std::collections::{BTreeMap, HashMap};
 
 /// The state multiproof of target accounts and multiproofs of their storage tries.
+/// Multiproof is effectively a state subtrie that only contains the nodes
+/// in the paths of target accounts.  
 #[derive(Clone, Default, Debug)]
 pub struct MultiProof {
     /// State trie multiproof for requested accounts.

crates/trie/db/src/proof.rs

@@ -4,9 +4,10 @@ use reth_execution_errors::StateProofError;
 use reth_primitives::{Address, B256};
 use reth_trie::{
     hashed_cursor::HashedPostStateCursorFactory, proof::Proof,
-    trie_cursor::InMemoryTrieCursorFactory, TrieInput,
+    trie_cursor::InMemoryTrieCursorFactory, MultiProof, TrieInput,
 };
 use reth_trie_common::AccountProof;
+use std::collections::{HashMap, HashSet};
 
 /// Extends [`Proof`] with operations specific for working with a database transaction.
 pub trait DatabaseProof<'a, TX> {
@@ -20,6 +21,13 @@ pub trait DatabaseProof<'a, TX> {
         address: Address,
         slots: &[B256],
     ) -> Result<AccountProof, StateProofError>;
+
+    /// Generates the state [`MultiProof`] for target hashed account and storage keys.
+    fn overlay_multiproof(
+        tx: &'a TX,
+        input: TrieInput,
+        targets: HashMap<B256, HashSet<B256>>,
+    ) -> Result<MultiProof, StateProofError>;
 }
 
 impl<'a, TX: DbTx> DatabaseProof<'a, TX>
@@ -50,4 +58,25 @@ impl<'a, TX: DbTx> DatabaseProof<'a, TX>
             .with_prefix_sets_mut(input.prefix_sets)
             .account_proof(address, slots)
     }
+
+    fn overlay_multiproof(
+        tx: &'a TX,
+        input: TrieInput,
+        targets: HashMap<B256, HashSet<B256>>,
+    ) -> Result<MultiProof, StateProofError> {
+        let nodes_sorted = input.nodes.into_sorted();
+        let state_sorted = input.state.into_sorted();
+        Self::from_tx(tx)
+            .with_trie_cursor_factory(InMemoryTrieCursorFactory::new(
+                DatabaseTrieCursorFactory::new(tx),
+                &nodes_sorted,
+            ))
+            .with_hashed_cursor_factory(HashedPostStateCursorFactory::new(
+                DatabaseHashedCursorFactory::new(tx),
+                &state_sorted,
+            ))
+            .with_prefix_sets_mut(input.prefix_sets)
+            .with_targets(targets)
+            .multiproof()
+    }
 }

crates/trie/trie/src/proof.rs

@@ -12,7 +12,7 @@ use reth_execution_errors::trie::StateProofError;
 use reth_trie_common::{
     proof::ProofRetainer, AccountProof, MultiProof, StorageMultiProof, TrieAccount,
 };
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 /// A struct for generating merkle proofs.
 ///
@@ -28,7 +28,7 @@ pub struct Proof<T, H> {
     /// A set of prefix sets that have changes.
     prefix_sets: TriePrefixSetsMut,
     /// Proof targets.
-    targets: HashMap<B256, Vec<B256>>,
+    targets: HashMap<B256, HashSet<B256>>,
 }
 
 impl<T, H> Proof<T, H> {
@@ -68,8 +68,13 @@ impl<T, H> Proof<T, H> {
         self
     }
 
+    /// Set the target account and slots.
+    pub fn with_target(self, target: (B256, HashSet<B256>)) -> Self {
+        self.with_targets(HashMap::from([target]))
+    }
+
     /// Set the target accounts and slots.
-    pub fn with_targets(mut self, targets: HashMap<B256, Vec<B256>>) -> Self {
+    pub fn with_targets(mut self, targets: HashMap<B256, HashSet<B256>>) -> Self {
         self.targets = targets;
         self
     }
@@ -87,10 +92,7 @@ where
         slots: &[B256],
     ) -> Result<AccountProof, StateProofError> {
         Ok(self
-            .with_targets(HashMap::from([(
-                keccak256(address),
-                slots.iter().map(keccak256).collect(),
-            )]))
+            .with_target((keccak256(address), slots.iter().map(keccak256).collect()))
             .multiproof()?
             .account_proof(address, slots)?)
     }

crates/trie/trie/src/witness.rs

@@ -79,11 +79,13 @@ where
         state: HashedPostState,
     ) -> Result<HashMap<B256, Bytes>, TrieWitnessError> {
         let proof_targets = HashMap::from_iter(
-            state.accounts.keys().map(|hashed_address| (*hashed_address, Vec::new())).chain(
-                state.storages.iter().map(|(hashed_address, storage)| {
+            state
+                .accounts
+                .keys()
+                .map(|hashed_address| (*hashed_address, HashSet::default()))
+                .chain(state.storages.iter().map(|(hashed_address, storage)| {
                     (*hashed_address, storage.storage.keys().copied().collect())
-                }),
-            ),
+                })),
         );
         let mut account_multiproof =
             Proof::new(self.trie_cursor_factory.clone(), self.hashed_cursor_factory.clone())
@@ -137,13 +139,13 @@ where
                 // Right pad the target with 0s.
                 let mut padded_key = key.pack();
                 padded_key.resize(32, 0);
-                let target = (hashed_address, Vec::from([B256::from_slice(&padded_key)]));
+                let target_key = B256::from_slice(&padded_key);
                 let mut proof = Proof::new(
                     self.trie_cursor_factory.clone(),
                     self.hashed_cursor_factory.clone(),
                 )
                 .with_prefix_sets_mut(self.prefix_sets.clone())
-                .with_targets(HashMap::from([target]))
+                .with_target((hashed_address, HashSet::from([target_key])))
                 .storage_multiproof(hashed_address)?;
 
                 // The subtree only contains the proof for a single target.
@@ -162,7 +164,7 @@ where
             let mut proof =
                 Proof::new(self.trie_cursor_factory.clone(), self.hashed_cursor_factory.clone())
                     .with_prefix_sets_mut(self.prefix_sets.clone())
-                    .with_targets(HashMap::from([(B256::from_slice(&padded_key), Vec::new())]))
+                    .with_target((B256::from_slice(&padded_key), HashSet::default()))
                     .multiproof()?;
 
             // The subtree only contains the proof for a single target.