diff --git a/crates/chain-state/src/in_memory.rs b/crates/chain-state/src/in_memory.rs index b1d9531b1c0d..d8e1355332ba 100644 --- a/crates/chain-state/src/in_memory.rs +++ b/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>, + ) -> ProviderResult { + Ok(MultiProof::default()) + } + fn witness( &self, _input: TrieInput, diff --git a/crates/chain-state/src/memory_overlay.rs b/crates/chain-state/src/memory_overlay.rs index fbd30a4d9262..9c4f5257e3d9 100644 --- a/crates/chain-state/src/memory_overlay.rs +++ b/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>, + ) -> ProviderResult { + let MemoryOverlayTrieState { nodes, state } = self.trie_state().clone(); + input.prepend_cached(nodes, state); + self.historical.multiproof(input, targets) + } + fn witness( &self, mut input: TrieInput, diff --git a/crates/revm/src/test_utils.rs b/crates/revm/src/test_utils.rs index e32e98e3ba47..fe377bc5fc3d 100644 --- a/crates/revm/src/test_utils.rs +++ b/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>, + ) -> ProviderResult { + unimplemented!("proof generation is not supported") + } + fn witness( &self, _input: TrieInput, diff --git a/crates/rpc/rpc-eth-types/src/cache/db.rs b/crates/rpc/rpc-eth-types/src/cache/db.rs index 1dcfae25fc1f..82f90ebc50ba 100644 --- a/crates/rpc/rpc-eth-types/src/cache/db.rs +++ b/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>, + ) -> ProviderResult { + self.0.multiproof(input, targets) + } + fn witness( &self, input: reth_trie::TrieInput, diff --git a/crates/storage/provider/src/providers/bundle_state_provider.rs b/crates/storage/provider/src/providers/bundle_state_provider.rs index 8a911d716e84..51824bb6ffd6 100644 --- a/crates/storage/provider/src/providers/bundle_state_provider.rs +++ b/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 StateProofProvider self.state_provider.proof(input, address, slots) } + fn multiproof( + &self, + mut input: reth_trie::TrieInput, + targets: HashMap>, + ) -> ProviderResult { + 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, diff --git a/crates/storage/provider/src/providers/state/historical.rs b/crates/storage/provider/src/providers/state/historical.rs index 6b4ee7fb8741..86f4800e1f87 100644 --- a/crates/storage/provider/src/providers/state/historical.rs +++ b/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::::into) } + fn multiproof( + &self, + mut input: TrieInput, + targets: HashMap>, + ) -> ProviderResult { + input.prepend(self.revert_state()?); + Proof::overlay_multiproof(self.tx, input, targets).map_err(Into::::into) + } + fn witness( &self, mut input: TrieInput, diff --git a/crates/storage/provider/src/providers/state/latest.rs b/crates/storage/provider/src/providers/state/latest.rs index 0e6d7f69e425..16ea5cdcd466 100644 --- a/crates/storage/provider/src/providers/state/latest.rs +++ b/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::::into) } + fn multiproof( + &self, + input: TrieInput, + targets: HashMap>, + ) -> ProviderResult { + Proof::overlay_multiproof(self.tx, input, targets).map_err(Into::::into) + } + fn witness( &self, input: TrieInput, diff --git a/crates/storage/provider/src/providers/state/macros.rs b/crates/storage/provider/src/providers/state/macros.rs index 161f89dd9c91..0e86f2fae5b6 100644 --- a/crates/storage/provider/src/providers/state/macros.rs +++ b/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; + fn multiproof(&self, input: reth_trie::TrieInput, targets: std::collections::HashMap>) -> reth_storage_errors::provider::ProviderResult; fn witness(&self, input: reth_trie::TrieInput, target: reth_trie::HashedPostState) -> reth_storage_errors::provider::ProviderResult>; } ); diff --git a/crates/storage/provider/src/test_utils/mock.rs b/crates/storage/provider/src/test_utils/mock.rs index b00aa2ff04f2..cf3e358b5bc0 100644 --- a/crates/storage/provider/src/test_utils/mock.rs +++ b/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>, + ) -> ProviderResult { + Ok(MultiProof::default()) + } + fn witness( &self, _input: TrieInput, diff --git a/crates/storage/provider/src/test_utils/noop.rs b/crates/storage/provider/src/test_utils/noop.rs index e5ef41a914df..0b04a74ed67b 100644 --- a/crates/storage/provider/src/test_utils/noop.rs +++ b/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>, + ) -> ProviderResult { + Ok(MultiProof::default()) + } + fn witness( &self, _input: TrieInput, diff --git a/crates/storage/storage-api/src/trie.rs b/crates/storage/storage-api/src/trie.rs index 5a3b74b30674..e41a15e107d3 100644 --- a/crates/storage/storage-api/src/trie.rs +++ b/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; + /// Generate [`MultiProof`] for target hashed account and corresponding + /// hashed storage slot keys. + fn multiproof( + &self, + input: TrieInput, + targets: HashMap>, + ) -> ProviderResult; + /// Get trie witness for provided state. fn witness( &self, diff --git a/crates/trie/common/src/proofs.rs b/crates/trie/common/src/proofs.rs index 87f9b4e4b6c2..df32b1cb9f6a 100644 --- a/crates/trie/common/src/proofs.rs +++ b/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. diff --git a/crates/trie/db/src/proof.rs b/crates/trie/db/src/proof.rs index cffe0495b29a..def14961afe4 100644 --- a/crates/trie/db/src/proof.rs +++ b/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; + + /// Generates the state [`MultiProof`] for target hashed account and storage keys. + fn overlay_multiproof( + tx: &'a TX, + input: TrieInput, + targets: HashMap>, + ) -> Result; } 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>, + ) -> Result { + 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() + } } diff --git a/crates/trie/trie/src/proof.rs b/crates/trie/trie/src/proof.rs index 0a8b23f9c702..8b9d2f9d09fb 100644 --- a/crates/trie/trie/src/proof.rs +++ b/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 { /// A set of prefix sets that have changes. prefix_sets: TriePrefixSetsMut, /// Proof targets. - targets: HashMap>, + targets: HashMap>, } impl Proof { @@ -68,8 +68,13 @@ impl Proof { self } + /// Set the target account and slots. + pub fn with_target(self, target: (B256, HashSet)) -> Self { + self.with_targets(HashMap::from([target])) + } + /// Set the target accounts and slots. - pub fn with_targets(mut self, targets: HashMap>) -> Self { + pub fn with_targets(mut self, targets: HashMap>) -> Self { self.targets = targets; self } @@ -87,10 +92,7 @@ where slots: &[B256], ) -> Result { 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)?) } diff --git a/crates/trie/trie/src/witness.rs b/crates/trie/trie/src/witness.rs index eb2862dd60ef..92f7d408c240 100644 --- a/crates/trie/trie/src/witness.rs +++ b/crates/trie/trie/src/witness.rs @@ -79,11 +79,13 @@ where state: HashedPostState, ) -> Result, 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.