Introduce HashedStorageProvider

crates/chain-state/src/in_memory.rs

@@ -949,12 +949,13 @@ mod tests {
     use reth_errors::ProviderResult;
     use reth_primitives::{Account, Bytecode, EthPrimitives, Receipt};
     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, StorageMultiProof, StorageProof, TrieInput,
     };
+    use revm::db::BundleAccount;
 
     fn create_mock_state(
         test_block_builder: &mut TestBlockBuilder<EthPrimitives>,
@@ -1053,6 +1054,12 @@ mod tests {
         }
     }
 
+    impl HashedStorageProvider for MockStateProvider {
+        fn hashed_storage(&self, _account: &BundleAccount) -> HashedStorage {
+            HashedStorage::default()
+        }
+    }
+
     impl StorageRootProvider for MockStateProvider {
         fn storage_root(
             &self,

crates/chain-state/src/memory_overlay.rs

@@ -8,14 +8,14 @@ use alloy_primitives::{
 use reth_errors::ProviderResult;
 use reth_primitives::{Account, Bytecode, NodePrimitives};
 use reth_storage_api::{
-    AccountReader, BlockHashReader, HashedPostStateProvider, StateProofProvider, StateProvider,
-    StateRootProvider, StorageRootProvider,
+    AccountReader, BlockHashReader, HashedPostStateProvider, HashedStorageProvider,
+    StateProofProvider, StateProvider, StateRootProvider, StorageRootProvider,
 };
 use reth_trie::{
     updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof,
     StorageMultiProof, TrieInput,
 };
-use revm::db::BundleState;
+use revm::db::{BundleAccount, BundleState};
 use std::sync::OnceLock;
 
 /// A state provider that stores references to in-memory blocks along with their state as well as a
@@ -225,6 +225,12 @@ macro_rules! impl_state_provider {
             }
         }
 
+        impl $($tokens)* HashedStorageProvider for $type {
+            fn hashed_storage(&self, account: &BundleAccount) -> HashedStorage {
+                self.historical.hashed_storage(account)
+            }
+        }
+
         impl $($tokens)* StateProvider for $type {
             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::{
@@ -156,6 +156,12 @@ impl HashedPostStateProvider for StateProviderTest {
     }
 }
 
+impl HashedStorageProvider for StateProviderTest {
+    fn hashed_storage(&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

@@ -8,7 +8,7 @@ use alloy_primitives::{
 };
 use reth_errors::ProviderResult;
 use reth_revm::{database::StateProviderDatabase, db::CacheDB, DatabaseRef};
-use reth_storage_api::{HashedPostStateProvider, StateProvider};
+use reth_storage_api::{HashedPostStateProvider, HashedStorageProvider, StateProvider};
 use reth_trie::HashedStorage;
 use revm::Database;
 
@@ -148,6 +148,12 @@ impl HashedPostStateProvider for StateProviderTraitObjWrapper<'_> {
     }
 }
 
+impl HashedStorageProvider for StateProviderTraitObjWrapper<'_> {
+    fn hashed_storage(&self, account: &revm::db::BundleAccount) -> HashedStorage {
+        self.0.hashed_storage(account)
+    }
+}
+
 impl StateProvider for StateProviderTraitObjWrapper<'_> {
     fn storage(
         &self,

crates/stages/stages/src/stages/execution.rs

@@ -262,8 +262,8 @@ where
         + StaticFileProviderFactory
         + StatsReader
         + BlockHashReader
-        + StateWriter<Receipt = reth_primitives::Receipt>
-        + StateCommitmentProvider,
+        + StateCommitmentProvider
+        + StateWriter<Receipt = reth_primitives::Receipt>,
 {
     /// Return the id of the stage
     fn id(&self) -> StageId {

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,
@@ -38,12 +40,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(account))
             .unwrap_or_default()
     }
 }
@@ -195,6 +192,14 @@ impl<SP: StateProvider, EDP: ExecutionDataProvider> HashedPostStateProvider
     }
 }
 
+impl<SP: StateProvider, EDP: ExecutionDataProvider> HashedStorageProvider
+    for BundleStateProvider<SP, EDP>
+{
+    fn hashed_storage(&self, account: &revm::db::BundleAccount) -> HashedStorage {
+        self.state_provider.hashed_storage(account)
+    }
+}
+
 impl<SP: StateProvider, EDP: ExecutionDataProvider> StateProvider for BundleStateProvider<SP, EDP> {
     fn storage(
         &self,

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

@@ -16,7 +16,8 @@ use reth_db_api::{
 };
 use reth_primitives::{Account, Bytecode};
 use reth_storage_api::{
-    BlockNumReader, DBProvider, StateCommitmentProvider, StateProofProvider, StorageRootProvider,
+    BlockNumReader, DBProvider, HashedStorageProvider, StateCommitmentProvider, StateProofProvider,
+    StorageRootProvider,
 };
 use reth_storage_errors::provider::ProviderResult;
 use reth_trie::{
@@ -155,7 +156,13 @@ impl<'b, Provider: DBProvider + BlockNumReader + StateCommitmentProvider>
             );
         }
 
-        Ok(HashedStorage::from_reverts(self.tx(), address, self.block_number)?)
+        Ok(
+            HashedStorage::from_reverts::<<Provider::StateCommitment as StateCommitment>::KeyHasher>(
+                self.tx(),
+                address,
+                self.block_number,
+            )?,
+        )
     }
 
     fn history_info<T, K>(
@@ -406,6 +413,16 @@ impl<Provider: StateCommitmentProvider> HashedPostStateProvider
     }
 }
 
+impl<Provider: StateCommitmentProvider + Send + Sync> HashedStorageProvider
+    for HistoricalStateProviderRef<'_, Provider>
+{
+    fn hashed_storage(&self, account: &revm::db::BundleAccount) -> HashedStorage {
+        HashedStorage::from_bundle_account::<
+            <Provider::StateCommitment as StateCommitment>::KeyHasher,
+        >(account)
+    }
+}
+
 impl<Provider: DBProvider + BlockNumReader + BlockHashReader + StateCommitmentProvider>
     StateProvider for HistoricalStateProviderRef<'_, Provider>
 {

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

@@ -10,7 +10,8 @@ use reth_db::tables;
 use reth_db_api::{cursor::DbDupCursorRO, transaction::DbTx};
 use reth_primitives::{Account, Bytecode};
 use reth_storage_api::{
-    DBProvider, StateCommitmentProvider, StateProofProvider, StorageRootProvider,
+    DBProvider, HashedStorageProvider, StateCommitmentProvider, StateProofProvider,
+    StorageRootProvider,
 };
 use reth_storage_errors::provider::{ProviderError, ProviderResult};
 use reth_trie::{
@@ -167,6 +168,16 @@ impl<Provider: DBProvider + StateCommitmentProvider> HashedPostStateProvider
     }
 }
 
+impl<Provider: StateCommitmentProvider + Send + Sync> HashedStorageProvider
+    for LatestStateProviderRef<'_, Provider>
+{
+    fn hashed_storage<'a>(&self, account: &revm::db::BundleAccount) -> HashedStorage {
+        HashedStorage::from_bundle_account::<
+            <Provider::StateCommitment as StateCommitment>::KeyHasher,
+        >(account)
+    }
+}
+
 impl<Provider: DBProvider + BlockHashReader + StateCommitmentProvider> StateProvider
     for LatestStateProviderRef<'_, Provider>
 {

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

@@ -60,6 +60,9 @@ macro_rules! delegate_provider_impls {
             HashedPostStateProvider $(where [$($generics)*])? {
                 fn hashed_post_state(&self, bundle_state: &revm::db::BundleState) -> reth_trie::HashedPostState;
             }
+            HashedStorageProvider $(where [$($generics)*])? {
+                fn hashed_storage(&self, bundle_state: &revm::db::BundleAccount) -> reth_trie::HashedStorage;
+            }
         );
     }
 }

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

@@ -29,7 +29,7 @@ use reth_primitives::{
 use reth_primitives_traits::SignedTransaction;
 use reth_stages_types::{StageCheckpoint, StageId};
 use reth_storage_api::{
-    DatabaseProviderFactory, HashedPostStateProvider, StageCheckpointReader,
+    DatabaseProviderFactory, HashedPostStateProvider, HashedStorageProvider, StageCheckpointReader,
     StateCommitmentProvider, StateProofProvider, StorageRootProvider,
 };
 use reth_storage_errors::provider::{ConsistentViewError, ProviderError, ProviderResult};
@@ -702,6 +702,12 @@ impl HashedPostStateProvider for MockEthProvider {
     }
 }
 
+impl HashedStorageProvider for MockEthProvider {
+    fn hashed_storage(&self, _account: &revm::db::BundleAccount) -> HashedStorage {
+        HashedStorage::default()
+    }
+}
+
 impl StateProvider for MockEthProvider {
     fn storage(
         &self,

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

@@ -28,7 +28,8 @@ use reth_primitives::{
 use reth_prune_types::{PruneCheckpoint, PruneSegment};
 use reth_stages_types::{StageCheckpoint, StageId};
 use reth_storage_api::{
-    HashedPostStateProvider, NodePrimitivesProvider, StateProofProvider, StorageRootProvider,
+    HashedPostStateProvider, HashedStorageProvider, NodePrimitivesProvider, StateProofProvider,
+    StorageRootProvider,
 };
 use reth_storage_errors::provider::ProviderResult;
 use reth_trie::{
@@ -419,6 +420,12 @@ impl HashedPostStateProvider for NoopProvider {
     }
 }
 
+impl HashedStorageProvider for NoopProvider {
+    fn hashed_storage(&self, _account: &revm::db::BundleAccount) -> HashedStorage {
+        HashedStorage::default()
+    }
+}
+
 impl StateProvider for NoopProvider {
     fn storage(
         &self,

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

@@ -1,6 +1,6 @@
 use super::{
-    AccountReader, BlockHashReader, BlockIdReader, StateProofProvider, StateRootProvider,
-    StorageRootProvider,
+    AccountReader, BlockHashReader, BlockIdReader, HashedStorageProvider, StateProofProvider,
+    StateRootProvider, StorageRootProvider,
 };
 use alloy_consensus::constants::KECCAK_EMPTY;
 use alloy_eips::{BlockId, BlockNumberOrTag};
@@ -24,6 +24,7 @@ pub trait StateProvider:
     + StorageRootProvider
     + StateProofProvider
     + HashedPostStateProvider
+    + HashedStorageProvider
     + Send
     + Sync
 {

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

@@ -2,6 +2,8 @@ use alloy_primitives::{Address, BlockNumber, B256};
 use reth_db_api::models::BlockNumberAddress;
 use reth_primitives::StorageEntry;
 use reth_storage_errors::provider::ProviderResult;
+use reth_trie::HashedStorage;
+use revm::db::BundleAccount;
 use std::{
     collections::{BTreeMap, BTreeSet},
     ops::RangeInclusive,
@@ -41,3 +43,10 @@ pub trait StorageChangeSetReader: Send + Sync {
         block_number: BlockNumber,
     ) -> ProviderResult<Vec<(BlockNumberAddress, StorageEntry)>>;
 }
+
+/// Provider of [`HashedStorage`]
+#[auto_impl::auto_impl(&, Arc, Box)]
+pub trait HashedStorageProvider: Send + Sync {
+    /// Construct [`HashedStorage`] from the provided [`BundleAccount`].
+    fn hashed_storage(&self, account: &BundleAccount) -> HashedStorage;
+}

crates/trie/db/src/storage.rs

@@ -6,7 +6,8 @@ use reth_db::{cursor::DbCursorRO, models::BlockNumberAddress, tables, DatabaseEr
 use reth_db_api::transaction::DbTx;
 use reth_execution_errors::StorageRootError;
 use reth_trie::{
-    hashed_cursor::HashedPostStateCursorFactory, HashedPostState, HashedStorage, StorageRoot,
+    hashed_cursor::HashedPostStateCursorFactory, HashedPostState, HashedStorage, KeyHasher,
+    StorageRoot,
 };
 
 #[cfg(feature = "metrics")]
@@ -32,7 +33,11 @@ pub trait DatabaseStorageRoot<'a, TX> {
 pub trait DatabaseHashedStorage<TX>: Sized {
     /// Initializes [`HashedStorage`] from reverts. Iterates over storage reverts from the specified
     /// block up to the current tip and aggregates them into hashed storage in reverse.
-    fn from_reverts(tx: &TX, address: Address, from: BlockNumber) -> Result<Self, DatabaseError>;
+    fn from_reverts<KH: KeyHasher>(
+        tx: &TX,
+        address: Address,
+        from: BlockNumber,
+    ) -> Result<Self, DatabaseError>;
 }
 
 impl<'a, TX: DbTx> DatabaseStorageRoot<'a, TX>
@@ -79,13 +84,17 @@ impl<'a, TX: DbTx> DatabaseStorageRoot<'a, TX>
 }
 
 impl<TX: DbTx> DatabaseHashedStorage<TX> for HashedStorage {
-    fn from_reverts(tx: &TX, address: Address, from: BlockNumber) -> Result<Self, DatabaseError> {
+    fn from_reverts<KH: KeyHasher>(
+        tx: &TX,
+        address: Address,
+        from: BlockNumber,
+    ) -> Result<Self, DatabaseError> {
         let mut storage = Self::new(false);
         let mut storage_changesets_cursor = tx.cursor_read::<tables::StorageChangeSets>()?;
         for entry in storage_changesets_cursor.walk_range(BlockNumberAddress((from, address))..)? {
             let (BlockNumberAddress((_, storage_address)), storage_change) = entry?;
             if storage_address == address {
-                let hashed_slot = keccak256(storage_change.key);
+                let hashed_slot = KH::hash_key(storage_change.key);
                 if let hash_map::Entry::Vacant(entry) = storage.storage.entry(hashed_slot) {
                     entry.insert(storage_change.value);
                 }