Introduce KeyHasherProvider

crates/chain-state/src/in_memory.rs

@@ -950,10 +950,12 @@ mod tests {
     use reth_primitives::{Account, Bytecode, EthPrimitives, Receipt};
     use reth_storage_api::{
         AccountReader, BlockHashReader, HashedPostStateProvider, HashedStorageProvider,
-        StateProofProvider, StateProvider, StateRootProvider, StorageRootProvider,
+        KeyHasherProvider, StateProofProvider, StateProvider, StateRootProvider,
+        StorageRootProvider,
     };
     use reth_trie::{
-        AccountProof, HashedStorage, MultiProof, StorageMultiProof, StorageProof, TrieInput,
+        AccountProof, HashedStorage, KeccakKeyHasher, KeyHasher, MultiProof, StorageMultiProof,
+        StorageProof, TrieInput,
     };
     use revm::db::BundleAccount;
 
@@ -1060,6 +1062,12 @@ mod tests {
         }
     }
 
+    impl KeyHasherProvider for MockStateProvider {
+        fn hash_key(&self, bytes: &[u8]) -> B256 {
+            KeccakKeyHasher::hash_key(bytes)
+        }
+    }
+
     impl StorageRootProvider for MockStateProvider {
         fn storage_root(
             &self,

crates/chain-state/src/memory_overlay.rs

@@ -1,15 +1,14 @@
 use super::ExecutedBlock;
 use alloy_consensus::BlockHeader;
 use alloy_primitives::{
-    keccak256,
     map::{HashMap, HashSet},
     Address, BlockNumber, Bytes, StorageKey, StorageValue, B256,
 };
 use reth_errors::ProviderResult;
 use reth_primitives::{Account, Bytecode, NodePrimitives};
 use reth_storage_api::{
     AccountReader, BlockHashReader, HashedPostStateProvider, HashedStorageProvider,
-    StateProofProvider, StateProvider, StateRootProvider, StorageRootProvider,
+    KeyHasherProvider, StateProofProvider, StateProvider, StateRootProvider, StorageRootProvider,
 };
 use reth_trie::{
     updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof,
@@ -152,7 +151,7 @@ macro_rules! impl_state_provider {
             fn storage_root(&self, address: Address, storage: HashedStorage) -> ProviderResult<B256> {
                 let state = &self.trie_state().state;
                 let mut hashed_storage =
-                    state.storages.get(&keccak256(address)).cloned().unwrap_or_default();
+                    state.storages.get(&self.hash_key(address.as_ref())).cloned().unwrap_or_default();
                 hashed_storage.extend(&storage);
                 self.historical.storage_root(address, hashed_storage)
             }
@@ -166,7 +165,7 @@ macro_rules! impl_state_provider {
             ) -> ProviderResult<reth_trie::StorageProof> {
                 let state = &self.trie_state().state;
                 let mut hashed_storage =
-                    state.storages.get(&keccak256(address)).cloned().unwrap_or_default();
+                    state.storages.get(&self.hash_key(address.as_ref())).cloned().unwrap_or_default();
                 hashed_storage.extend(&storage);
                 self.historical.storage_proof(address, slot, hashed_storage)
             }
@@ -180,7 +179,7 @@ macro_rules! impl_state_provider {
             ) -> ProviderResult<StorageMultiProof> {
                 let state = &self.trie_state().state;
                 let mut hashed_storage =
-                    state.storages.get(&keccak256(address)).cloned().unwrap_or_default();
+                    state.storages.get(&self.hash_key(address.as_ref())).cloned().unwrap_or_default();
                 hashed_storage.extend(&storage);
                 self.historical.storage_multiproof(address, slots, hashed_storage)
             }
@@ -231,6 +230,12 @@ macro_rules! impl_state_provider {
             }
         }
 
+        impl $($tokens)* KeyHasherProvider for $type {
+            fn hash_key(&self, bytes: &[u8]) -> B256 {
+                self.historical.hash_key(bytes)
+            }
+        }
+
         impl $($tokens)* StateProvider for $type {
             fn storage(
                 &self,

crates/engine/invalid-block-hooks/src/witness.rs

@@ -1,5 +1,5 @@
 use alloy_consensus::Header;
-use alloy_primitives::{keccak256, B256, U256};
+use alloy_primitives::{B256, U256};
 use alloy_rpc_types_debug::ExecutionWitness;
 use eyre::OptionExt;
 use pretty_assertions::Comparison;
@@ -69,9 +69,9 @@ where
         // Setup database.
         let provider = self.provider.state_by_block_hash(parent_header.hash())?;
         #[cfg(not(feature = "scroll"))]
-        let state = reth_revm::database::StateProviderDatabase::new(provider);
+        let state = reth_revm::database::StateProviderDatabase::new(&provider);
         #[cfg(feature = "scroll")]
-        let state = reth_scroll_storage::ScrollStateProviderDatabase::new(provider);
+        let state = reth_scroll_storage::ScrollStateProviderDatabase::new(&provider);
         let mut db = StateBuilder::new().with_database(state).with_bundle_update().build();
 
         // Setup environment for the execution.
@@ -129,7 +129,7 @@ where
         // referenced accounts + storage slots.
         let mut hashed_state = db.database.hashed_post_state(&bundle_state);
         for (address, account) in db.cache.accounts {
-            let hashed_address = keccak256(address);
+            let hashed_address = provider.hash_key(address.as_ref());
             #[cfg(feature = "scroll")]
             let hashed_account = account.account.as_ref().map(|a| {
                 Into::<reth_scroll_revm::AccountInfo>::into((
@@ -152,7 +152,7 @@ where
 
                 for (slot, value) in account.storage {
                     let slot = B256::from(slot);
-                    let hashed_slot = keccak256(slot);
+                    let hashed_slot = provider.hash_key(slot.as_ref());
                     storage.storage.insert(hashed_slot, value);
 
                     state_preimages.insert(hashed_slot, alloy_rlp::encode(slot).into());

crates/revm/src/test_utils.rs

@@ -7,11 +7,11 @@ use alloy_primitives::{
 use reth_primitives::{Account, Bytecode};
 use reth_storage_api::{
     AccountReader, BlockHashReader, HashedPostStateProvider, HashedStorageProvider,
-    StateProofProvider, StateProvider, StateRootProvider, StorageRootProvider,
+    KeyHasherProvider, StateProofProvider, StateProvider, StateRootProvider, StorageRootProvider,
 };
 use reth_storage_errors::provider::ProviderResult;
 use reth_trie::{
-    updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, KeccakKeyHasher,
+    updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, KeccakKeyHasher, KeyHasher,
     MultiProof, StorageMultiProof, StorageProof, TrieInput,
 };
 
@@ -162,6 +162,12 @@ impl HashedStorageProvider for StateProviderTest {
     }
 }
 
+impl KeyHasherProvider for StateProviderTest {
+    fn hash_key(&self, bytes: &[u8]) -> B256 {
+        KeccakKeyHasher::hash_key(bytes)
+    }
+}
+
 impl StateProvider for StateProviderTest {
     fn storage(
         &self,

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

@@ -8,7 +8,9 @@ use alloy_primitives::{
 };
 use reth_errors::ProviderResult;
 use reth_revm::{database::StateProviderDatabase, db::CacheDB, DatabaseRef};
-use reth_storage_api::{HashedPostStateProvider, HashedStorageProvider, StateProvider};
+use reth_storage_api::{
+    HashedPostStateProvider, HashedStorageProvider, KeyHasherProvider, StateProvider,
+};
 use reth_trie::HashedStorage;
 use revm::Database;
 
@@ -154,6 +156,12 @@ impl HashedStorageProvider for StateProviderTraitObjWrapper<'_> {
     }
 }
 
+impl KeyHasherProvider for StateProviderTraitObjWrapper<'_> {
+    fn hash_key(&self, bytes: &[u8]) -> B256 {
+        self.0.hash_key(bytes)
+    }
+}
+
 impl StateProvider for StateProviderTraitObjWrapper<'_> {
     fn storage(
         &self,

crates/stages/stages/src/stages/hashing_account.rs

@@ -1,4 +1,4 @@
-use alloy_primitives::{keccak256, B256};
+use alloy_primitives::B256;
 use itertools::Itertools;
 use reth_config::config::{EtlConfig, HashingConfig};
 use reth_db::{tables, RawKey, RawTable, RawValue};
@@ -8,12 +8,16 @@ use reth_db_api::{
 };
 use reth_etl::Collector;
 use reth_primitives::Account;
-use reth_provider::{AccountExtReader, DBProvider, HashingWriter, StatsReader};
+use reth_provider::{
+    AccountExtReader, DBProvider, HashingWriter, StateCommitmentProvider, StatsReader,
+};
 use reth_stages_api::{
     AccountHashingCheckpoint, EntitiesCheckpoint, ExecInput, ExecOutput, Stage, StageCheckpoint,
     StageError, StageId, UnwindInput, UnwindOutput,
 };
 use reth_storage_errors::provider::ProviderResult;
+use reth_trie::KeyHasher;
+use reth_trie_db::StateCommitment;
 use std::{
     fmt::Debug,
     ops::{Range, RangeInclusive},
@@ -133,7 +137,11 @@ impl Default for AccountHashingStage {
 
 impl<Provider> Stage<Provider> for AccountHashingStage
 where
-    Provider: DBProvider<Tx: DbTxMut> + HashingWriter + AccountExtReader + StatsReader,
+    Provider: DBProvider<Tx: DbTxMut>
+        + HashingWriter
+        + AccountExtReader
+        + StatsReader
+        + StateCommitmentProvider,
 {
     /// Return the id of the stage
     fn id(&self) -> StageId {
@@ -174,7 +182,14 @@ where
                 rayon::spawn(move || {
                     for (address, account) in chunk {
                         let address = address.key().unwrap();
-                        let _ = tx.send((RawKey::new(keccak256(address)), account));
+                        let _ = tx.send((
+                            RawKey::new(
+                                <<Provider::StateCommitment as StateCommitment>::KeyHasher as KeyHasher>::hash_key(
+                                    address,
+                                ),
+                            ),
+                            account,
+                        ));
                     }
                 });
 
@@ -301,7 +316,7 @@ mod tests {
         stage_test_suite_ext, ExecuteStageTestRunner, StageTestRunner, TestRunnerError,
         UnwindStageTestRunner,
     };
-    use alloy_primitives::U256;
+    use alloy_primitives::{keccak256, U256};
     use assert_matches::assert_matches;
     use reth_primitives::Account;
     use reth_provider::providers::StaticFileWriter;

crates/stages/stages/src/stages/hashing_storage.rs

@@ -1,4 +1,4 @@
-use alloy_primitives::{bytes::BufMut, keccak256, B256};
+use alloy_primitives::{bytes::BufMut, B256};
 use itertools::Itertools;
 use reth_config::config::{EtlConfig, HashingConfig};
 use reth_db::tables;
@@ -10,12 +10,16 @@ use reth_db_api::{
 };
 use reth_etl::Collector;
 use reth_primitives::StorageEntry;
-use reth_provider::{DBProvider, HashingWriter, StatsReader, StorageReader};
+use reth_provider::{
+    DBProvider, HashingWriter, StateCommitmentProvider, StatsReader, StorageReader,
+};
 use reth_stages_api::{
     EntitiesCheckpoint, ExecInput, ExecOutput, Stage, StageCheckpoint, StageError, StageId,
     StorageHashingCheckpoint, UnwindInput, UnwindOutput,
 };
 use reth_storage_errors::provider::ProviderResult;
+use reth_trie::KeyHasher;
+use reth_trie_db::StateCommitment;
 use std::{
     fmt::Debug,
     sync::mpsc::{self, Receiver},
@@ -64,7 +68,11 @@ impl Default for StorageHashingStage {
 
 impl<Provider> Stage<Provider> for StorageHashingStage
 where
-    Provider: DBProvider<Tx: DbTxMut> + StorageReader + HashingWriter + StatsReader,
+    Provider: DBProvider<Tx: DbTxMut>
+        + StorageReader
+        + HashingWriter
+        + StatsReader
+        + StateCommitmentProvider,
 {
     /// Return the id of the stage
     fn id(&self) -> StageId {
@@ -103,8 +111,13 @@ where
                 rayon::spawn(move || {
                     for (address, slot) in chunk {
                         let mut addr_key = Vec::with_capacity(64);
-                        addr_key.put_slice(keccak256(address).as_slice());
-                        addr_key.put_slice(keccak256(slot.key).as_slice());
+                        addr_key.put_slice(
+                            <<Provider::StateCommitment as StateCommitment>::KeyHasher as KeyHasher>::hash_key(
+                                address,
+                            )
+                            .as_slice(),
+                        );
+                        addr_key.put_slice(<<Provider::StateCommitment as StateCommitment>::KeyHasher as KeyHasher>::hash_key(slot.key).as_slice());
                         let _ = tx.send((addr_key, CompactU256::from(slot.value)));
                     }
                 });
@@ -212,7 +225,7 @@ mod tests {
         stage_test_suite_ext, ExecuteStageTestRunner, StageTestRunner, TestRunnerError,
         TestStageDB, UnwindStageTestRunner,
     };
-    use alloy_primitives::{Address, U256};
+    use alloy_primitives::{keccak256, Address, U256};
     use assert_matches::assert_matches;
     use rand::Rng;
     use reth_db_api::{

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

@@ -7,7 +7,8 @@ use alloy_primitives::{
 };
 use reth_primitives::{Account, Bytecode};
 use reth_storage_api::{
-    HashedPostStateProvider, HashedStorageProvider, StateProofProvider, StorageRootProvider,
+    HashedPostStateProvider, HashedStorageProvider, KeyHasherProvider, StateProofProvider,
+    StorageRootProvider,
 };
 use reth_storage_errors::provider::ProviderResult;
 use reth_trie::{
@@ -200,6 +201,14 @@ impl<SP: StateProvider, EDP: ExecutionDataProvider> HashedStorageProvider
     }
 }
 
+impl<SP: StateProvider, EDP: ExecutionDataProvider> KeyHasherProvider
+    for BundleStateProvider<SP, EDP>
+{
+    fn hash_key(&self, bytes: &[u8]) -> B256 {
+        self.state_provider.hash_key(bytes)
+    }
+}
+
 impl<SP: StateProvider, EDP: ExecutionDataProvider> StateProvider for BundleStateProvider<SP, EDP> {
     fn storage(
         &self,