feat: introduce HashedStorageProvider

crates/chain-state/src/in_memory.rs

@@ -870,8 +870,8 @@ mod tests {
     use reth_errors::ProviderResult;
     use reth_primitives::{Account, Bytecode, Receipt, Requests};
     use reth_storage_api::{
-        AccountReader, BlockHashReader, HashedPostStateProvider, StateProofProvider, StateProvider,
-        StateRootProvider, StorageRootProvider,
+        AccountReader, BlockHashReader, HashedPostStateProvider, HashedStorageProvider,
+        StateProofProvider, StateProvider, StateRootProvider, StorageRootProvider,
     };
     use reth_trie::{AccountProof, HashedStorage, MultiProof, StorageProof, TrieInput};
 
@@ -1028,6 +1028,15 @@ mod tests {
         }
     }
 
+    impl HashedStorageProvider for MockStateProvider {
+        fn hashed_storage_from_bundle_account(
+            &self,
+            _account: &reth_execution_types::BundleAccount,
+        ) -> HashedStorage {
+            HashedStorage::default()
+        }
+    }
+
     #[test]
     fn test_in_memory_state_impl_state_by_hash() {
         let mut state_by_hash = HashMap::default();

crates/chain-state/src/memory_overlay.rs

@@ -7,8 +7,8 @@ use alloy_primitives::{
 use reth_errors::ProviderResult;
 use reth_primitives::{Account, Bytecode};
 use reth_storage_api::{
-    AccountReader, BlockHashReader, HashedPostStateProvider, StateProofProvider, StateProvider,
-    StateProviderBox, StateRootProvider, StorageRootProvider,
+    AccountReader, BlockHashReader, HashedPostStateProvider, HashedStorageProvider,
+    StateProofProvider, StateProvider, StateProviderBox, StateRootProvider, StorageRootProvider,
 };
 use reth_trie::{
     updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof, TrieInput,
@@ -204,6 +204,15 @@ impl HashedPostStateProvider for MemoryOverlayStateProvider {
     }
 }
 
+impl HashedStorageProvider for MemoryOverlayStateProvider {
+    fn hashed_storage_from_bundle_account(
+        &self,
+        account: &reth_execution_types::BundleAccount,
+    ) -> HashedStorage {
+        self.historical.hashed_storage_from_bundle_account(account)
+    }
+}
+
 impl StateProvider for MemoryOverlayStateProvider {
     fn storage(
         &self,

crates/revm/src/test_utils.rs

@@ -6,8 +6,8 @@ use alloy_primitives::{
 };
 use reth_primitives::{Account, Bytecode};
 use reth_storage_api::{
-    AccountReader, BlockHashReader, HashedPostStateProvider, StateProofProvider, StateProvider,
-    StateRootProvider, StorageRootProvider,
+    AccountReader, BlockHashReader, HashedPostStateProvider, HashedStorageProvider,
+    StateProofProvider, StateProvider, StateRootProvider, StorageRootProvider,
 };
 use reth_storage_errors::provider::ProviderResult;
 use reth_trie::{
@@ -155,6 +155,15 @@ impl HashedPostStateProvider for StateProviderTest {
     }
 }
 
+impl HashedStorageProvider for StateProviderTest {
+    fn hashed_storage_from_bundle_account(
+        &self,
+        account: &revm::db::BundleAccount,
+    ) -> HashedStorage {
+        HashedStorage::from_bundle_account::<KeccakKeyHasher>(account)
+    }
+}
+
 impl StateProvider for StateProviderTest {
     fn storage(
         &self,

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

@@ -146,6 +146,15 @@ impl reth_storage_api::HashedPostStateProvider for StateProviderTraitObjWrapper<
     }
 }
 
+impl reth_storage_api::HashedStorageProvider for StateProviderTraitObjWrapper<'_> {
+    fn hashed_storage_from_bundle_account(
+        &self,
+        account: &reth_execution_types::BundleAccount,
+    ) -> reth_trie::HashedStorage {
+        self.0.hashed_storage_from_bundle_account(account)
+    }
+}
+
 impl StateProvider for StateProviderTraitObjWrapper<'_> {
     fn account_balance(
         &self,

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

@@ -6,7 +6,9 @@ use alloy_primitives::{
     Address, BlockNumber, Bytes, B256,
 };
 use reth_primitives::{Account, Bytecode};
-use reth_storage_api::{HashedPostStateProvider, StateProofProvider, StorageRootProvider};
+use reth_storage_api::{
+    HashedPostStateProvider, HashedStorageProvider, StateProofProvider, StorageRootProvider,
+};
 use reth_storage_errors::provider::ProviderResult;
 use reth_trie::{
     updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof, TrieInput,
@@ -37,12 +39,7 @@ impl<SP: StateProvider, EDP: ExecutionDataProvider> BundleStateProvider<SP, EDP>
         let bundle_state = self.block_execution_data_provider.execution_outcome().state();
         bundle_state
             .account(&address)
-            .map(|account| {
-                HashedStorage::from_plain_storage(
-                    account.status,
-                    account.storage.iter().map(|(slot, value)| (slot, &value.present_value)),
-                )
-            })
+            .map(|account| self.state_provider.hashed_storage_from_bundle_account(account))
             .unwrap_or_default()
     }
 }
@@ -193,6 +190,17 @@ impl<SP: StateProvider, EDP: ExecutionDataProvider> HashedPostStateProvider
     }
 }
 
+impl<SP: StateProvider, EDP: ExecutionDataProvider> HashedStorageProvider
+    for BundleStateProvider<SP, EDP>
+{
+    fn hashed_storage_from_bundle_account(
+        &self,
+        account: &reth_execution_types::BundleAccount,
+    ) -> HashedStorage {
+        self.state_provider.hashed_storage_from_bundle_account(account)
+    }
+}
+
 impl<SP: StateProvider, EDP: ExecutionDataProvider> StateProvider for BundleStateProvider<SP, EDP> {
     fn storage(
         &self,

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

@@ -14,7 +14,9 @@ use reth_db_api::{
     transaction::DbTx,
 };
 use reth_primitives::{constants::EPOCH_SLOTS, Account, Bytecode, StaticFileSegment};
-use reth_storage_api::{HashedPostStateProvider, StateProofProvider, StorageRootProvider};
+use reth_storage_api::{
+    HashedPostStateProvider, HashedStorageProvider, StateProofProvider, StorageRootProvider,
+};
 use reth_storage_errors::provider::ProviderResult;
 use reth_trie::{
     proof::{Proof, StorageProof},
@@ -428,6 +430,17 @@ impl<TX: DbTx, SC: StateCommitment> HashedPostStateProvider
     }
 }
 
+impl<TX: DbTx, SC: StateCommitment> HashedStorageProvider
+    for HistoricalStateProviderRef<'_, TX, SC>
+{
+    fn hashed_storage_from_bundle_account(
+        &self,
+        account: &reth_execution_types::BundleAccount,
+    ) -> HashedStorage {
+        HashedStorage::from_bundle_account::<SC::KeyHasher>(account)
+    }
+}
+
 impl<TX: DbTx, SC: StateCommitment> StateProvider for HistoricalStateProviderRef<'_, TX, SC> {
     /// Get storage.
     fn storage(

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

@@ -12,7 +12,9 @@ use reth_db_api::{
     transaction::DbTx,
 };
 use reth_primitives::{Account, Bytecode, StaticFileSegment};
-use reth_storage_api::{HashedPostStateProvider, StateProofProvider, StorageRootProvider};
+use reth_storage_api::{
+    HashedPostStateProvider, HashedStorageProvider, StateProofProvider, StorageRootProvider,
+};
 use reth_storage_errors::provider::{ProviderError, ProviderResult};
 use reth_trie::{
     proof::{Proof, StorageProof},
@@ -180,6 +182,15 @@ impl<TX: DbTx, SC: StateCommitment> HashedPostStateProvider for LatestStateProvi
     }
 }
 
+impl<TX: DbTx, SC: StateCommitment> HashedStorageProvider for LatestStateProviderRef<'_, TX, SC> {
+    fn hashed_storage_from_bundle_account(
+        &self,
+        account: &reth_execution_types::BundleAccount,
+    ) -> HashedStorage {
+        HashedStorage::from_bundle_account::<SC::KeyHasher>(account)
+    }
+}
+
 impl<TX: DbTx, SC: StateCommitment> StateProvider for LatestStateProviderRef<'_, TX, SC> {
     /// Get storage.
     fn storage(

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

@@ -60,6 +60,9 @@ macro_rules! delegate_provider_impls {
                 fn hashed_post_state_from_reverts(&self, block_number: alloy_primitives::BlockNumber) -> reth_storage_errors::provider::ProviderResult<reth_trie::HashedPostState>;
                 fn hashed_post_state_from_bundle_state(&self, bundle_state: &revm::db::BundleState) -> reth_trie::HashedPostState;
             }
+            HashedStorageProvider $(where [$($generics)*])? {
+                fn hashed_storage_from_bundle_account(&self, account: &revm::db::BundleAccount) -> reth_trie::HashedStorage;
+            }
         );
     }
 }

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

@@ -27,8 +27,8 @@ use reth_primitives::{
 };
 use reth_stages_types::{StageCheckpoint, StageId};
 use reth_storage_api::{
-    DatabaseProviderFactory, HashedPostStateProvider, StageCheckpointReader, StateProofProvider,
-    StorageRootProvider,
+    DatabaseProviderFactory, HashedPostStateProvider, HashedStorageProvider, StageCheckpointReader,
+    StateProofProvider, StorageRootProvider,
 };
 use reth_storage_errors::provider::{ConsistentViewError, ProviderError, ProviderResult};
 use reth_trie::{
@@ -687,6 +687,15 @@ impl HashedPostStateProvider for MockEthProvider {
     }
 }
 
+impl HashedStorageProvider for MockEthProvider {
+    fn hashed_storage_from_bundle_account(
+        &self,
+        account: &reth_execution_types::BundleAccount,
+    ) -> HashedStorage {
+        HashedStorage::from_bundle_account::<KeccakKeyHasher>(account)
+    }
+}
+
 impl StateProvider for MockEthProvider {
     fn storage(
         &self,

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

@@ -24,7 +24,9 @@ use reth_primitives::{
 };
 use reth_prune_types::{PruneCheckpoint, PruneSegment};
 use reth_stages_types::{StageCheckpoint, StageId};
-use reth_storage_api::{HashedPostStateProvider, StateProofProvider, StorageRootProvider};
+use reth_storage_api::{
+    HashedPostStateProvider, HashedStorageProvider, StateProofProvider, StorageRootProvider,
+};
 use reth_storage_errors::provider::ProviderResult;
 use reth_trie::{
     updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof, TrieInput,
@@ -410,6 +412,15 @@ impl HashedPostStateProvider for NoopProvider {
     }
 }
 
+impl HashedStorageProvider for NoopProvider {
+    fn hashed_storage_from_bundle_account(
+        &self,
+        _account: &reth_execution_types::BundleAccount,
+    ) -> HashedStorage {
+        HashedStorage::default()
+    }
+}
+
 impl StateProvider for NoopProvider {
     fn storage(
         &self,

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

@@ -1,3 +1,5 @@
+use crate::HashedStorageProvider;
+
 use super::{
     AccountReader, BlockHashReader, BlockIdReader, StateProofProvider, StateRootProvider,
     StorageRootProvider,
@@ -23,6 +25,7 @@ pub trait StateProvider:
     + StorageRootProvider
     + StateProofProvider
     + HashedPostStateProvider
+    + HashedStorageProvider
     + Send
     + Sync
 {

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

@@ -1,7 +1,9 @@
 use alloy_primitives::{Address, BlockNumber, B256};
 use reth_db_api::models::BlockNumberAddress;
+use reth_execution_types::BundleAccount;
 use reth_primitives::StorageEntry;
 use reth_storage_errors::provider::ProviderResult;
+use reth_trie::HashedStorage;
 use std::{
     collections::{BTreeMap, BTreeSet},
     ops::RangeInclusive,
@@ -32,6 +34,13 @@ pub trait StorageReader: Send + Sync {
     ) -> ProviderResult<BTreeMap<(Address, B256), Vec<u64>>>;
 }
 
+/// [`HashedStorage`] provider
+#[auto_impl::auto_impl(&, Arc, Box)]
+pub trait HashedStorageProvider: Send + Sync {
+    /// Returns the [`HashedStorage`] for the provided `BundleAccount`.
+    fn hashed_storage_from_bundle_account(&self, account: &BundleAccount) -> HashedStorage;
+}
+
 /// Storage ChangeSet reader
 #[auto_impl::auto_impl(&, Arc, Box)]
 pub trait StorageChangeSetReader: Send + Sync {

crates/trie/trie/src/state.rs

@@ -2,7 +2,7 @@ use crate::{
     prefix_set::{PrefixSetMut, TriePrefixSetsMut},
     Nibbles,
 };
-use alloy_primitives::{keccak256, Address, B256, U256};
+use alloy_primitives::{Address, B256, U256};
 use itertools::Itertools;
 use rayon::prelude::{IntoParallelIterator, ParallelIterator};
 use reth_primitives::Account;
@@ -34,7 +34,7 @@ impl HashedPostState {
             .map(|(address, account)| {
                 let hashed_address = KH::hash_key(address);
                 let hashed_account = account.info.clone().map(Into::into);
-                let hashed_storage = HashedStorage::from_plain_storage(
+                let hashed_storage = HashedStorage::from_plain_storage::<KH>(
                     account.status,
                     account.storage.iter().map(|(slot, value)| (slot, &value.present_value)),
                 );
@@ -61,7 +61,7 @@ impl HashedPostState {
             .map(|(address, account)| {
                 let hashed_address = KH::hash_key(address);
                 let hashed_account = account.account.as_ref().map(|a| a.info.clone().into());
-                let hashed_storage = HashedStorage::from_plain_storage(
+                let hashed_storage = HashedStorage::from_plain_storage::<KH>(
                     account.status,
                     account.account.as_ref().map(|a| a.storage.iter()).into_iter().flatten(),
                 );
@@ -220,13 +220,24 @@ impl HashedStorage {
     }
 
     /// Create new hashed storage from account status and plain storage.
-    pub fn from_plain_storage<'a>(
+    pub fn from_plain_storage<'a, KH: KeyHasher>(
         status: AccountStatus,
         storage: impl IntoIterator<Item = (&'a U256, &'a U256)>,
     ) -> Self {
         Self::from_iter(
             status.was_destroyed(),
-            storage.into_iter().map(|(key, value)| (keccak256(B256::from(*key)), *value)),
+            storage.into_iter().map(|(key, value)| (KH::hash_key(B256::from(*key)), *value)),
+        )
+    }
+
+    /// Create a new hashed storage from the provided [`BundleAccount`]
+    ///
+    /// This function will use the present value of the storage slots in the account to create the
+    /// hashed storage.
+    pub fn from_bundle_account<KH: KeyHasher>(account: &BundleAccount) -> Self {
+        Self::from_plain_storage::<KH>(
+            account.status,
+            account.storage.iter().map(|(slot, value)| (slot, &value.present_value)),
         )
     }