feat : add storage_root provider function for account (#9659)

Co-authored-by: Roman Krasiuk <rokrassyuk@gmail.com>

crates/chain-state/src/in_memory.rs

@@ -673,7 +673,7 @@ mod tests {
     use reth_storage_api::{
         AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateRootProvider,
     };
-    use reth_trie::AccountProof;
+    use reth_trie::{AccountProof, HashedStorage};
 
     fn create_mock_state(
         test_block_builder: &mut TestBlockBuilder,
@@ -753,6 +753,14 @@ mod tests {
         ) -> ProviderResult<(B256, TrieUpdates)> {
             Ok((B256::random(), TrieUpdates::default()))
         }
+
+        fn hashed_storage_root(
+            &self,
+            _address: Address,
+            _hashed_storage: HashedStorage,
+        ) -> ProviderResult<B256> {
+            Ok(B256::random())
+        }
     }
 
     impl StateProofProvider for MockStateProvider {

crates/chain-state/src/memory_overlay.rs

@@ -8,7 +8,7 @@ use reth_primitives::{
 use reth_storage_api::{
     AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateRootProvider,
 };
-use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState};
+use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage};
 
 /// A state provider that stores references to in-memory blocks along with their state as well as
 /// the historical state provider for fallback lookups.
@@ -101,6 +101,15 @@ impl StateRootProvider for MemoryOverlayStateProvider {
         state.extend(hashed_state);
         self.historical.hashed_state_root_with_updates(state)
     }
+
+    // TODO: Currently this does not reuse available in-memory trie nodes.
+    fn hashed_storage_root(
+        &self,
+        address: Address,
+        storage: HashedStorage,
+    ) -> ProviderResult<B256> {
+        self.historical.hashed_storage_root(address, storage)
+    }
 }
 
 impl StateProofProvider for MemoryOverlayStateProvider {

crates/evm/execution-errors/src/trie.rs

@@ -33,6 +33,14 @@ pub enum StorageRootError {
     Database(#[from] DatabaseError),
 }
 
+impl From<StorageRootError> for DatabaseError {
+    fn from(err: StorageRootError) -> Self {
+        match err {
+            StorageRootError::Database(err) => err,
+        }
+    }
+}
+
 /// State proof errors.
 #[derive(Error, Debug, PartialEq, Eq, Clone)]
 pub enum StateProofError {

crates/revm/src/test_utils.rs

@@ -6,7 +6,7 @@ use reth_storage_api::{
     AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateRootProvider,
 };
 use reth_storage_errors::provider::ProviderResult;
-use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState};
+use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage};
 
 #[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
@@ -78,6 +78,14 @@ impl StateRootProvider for StateProviderTest {
     ) -> ProviderResult<(B256, TrieUpdates)> {
         unimplemented!("state root computation is not supported")
     }
+
+    fn hashed_storage_root(
+        &self,
+        _address: Address,
+        _hashed_storage: HashedStorage,
+    ) -> ProviderResult<B256> {
+        unimplemented!("storage root is not supported")
+    }
 }
 
 impl StateProofProvider for StateProviderTest {

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

@@ -2,9 +2,11 @@
 //! <https://github.com/rust-lang/rust/issues/100013> in default implementation of
 //! `reth_rpc_eth_api::helpers::Call`.
 
-use reth_primitives::{B256, U256};
+use reth_errors::ProviderResult;
+use reth_primitives::{Address, B256, U256};
 use reth_provider::StateProvider;
 use reth_revm::{database::StateProviderDatabase, db::CacheDB, DatabaseRef};
+use reth_trie::HashedStorage;
 use revm::Database;
 
 /// Helper alias type for the state's [`CacheDB`]
@@ -29,6 +31,14 @@ impl<'a> reth_provider::StateRootProvider for StateProviderTraitObjWrapper<'a> {
     ) -> reth_errors::ProviderResult<(B256, reth_trie::updates::TrieUpdates)> {
         self.0.hashed_state_root_with_updates(hashed_state)
     }
+
+    fn hashed_storage_root(
+        &self,
+        address: Address,
+        hashed_storage: HashedStorage,
+    ) -> ProviderResult<B256> {
+        self.0.hashed_storage_root(address, hashed_storage)
+    }
 }
 
 impl<'a> reth_provider::StateProofProvider for StateProviderTraitObjWrapper<'a> {

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

@@ -6,7 +6,7 @@ use crate::{
 use reth_primitives::{Account, Address, BlockNumber, Bytecode, Bytes, B256};
 use reth_storage_api::StateProofProvider;
 use reth_storage_errors::provider::ProviderResult;
-use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState};
+use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage};
 use revm::db::BundleState;
 
 /// A state provider that resolves to data from either a wrapped [`crate::ExecutionOutcome`]
@@ -98,6 +98,20 @@ impl<SP: StateProvider, EDP: ExecutionDataProvider> StateRootProvider
         state.extend(hashed_state);
         self.state_provider.hashed_state_root_with_updates(state)
     }
+
+    fn hashed_storage_root(
+        &self,
+        address: Address,
+        hashed_storage: HashedStorage,
+    ) -> ProviderResult<B256> {
+        let bundle_state = self.block_execution_data_provider.execution_outcome().state();
+        let mut storage = bundle_state
+            .account(&address)
+            .map(|account| HashedStorage::from_bundle_state(account.status, &account.storage))
+            .unwrap_or_else(|| HashedStorage::new(false));
+        storage.extend(hashed_storage);
+        self.state_provider.hashed_storage_root(address, storage)
+    }
 }
 
 impl<SP: StateProvider, EDP: ExecutionDataProvider> StateProofProvider

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

@@ -17,10 +17,11 @@ use reth_storage_api::StateProofProvider;
 use reth_storage_errors::provider::ProviderResult;
 use reth_trie::{
     proof::Proof, updates::TrieUpdates, witness::TrieWitness, AccountProof, HashedPostState,
-    StateRoot,
+    HashedStorage, StateRoot, StorageRoot,
 };
 use reth_trie_db::{
-    DatabaseHashedPostState, DatabaseProof, DatabaseStateRoot, DatabaseTrieWitness,
+    DatabaseHashedPostState, DatabaseProof, DatabaseStateRoot, DatabaseStorageRoot,
+    DatabaseTrieWitness,
 };
 use std::{collections::HashMap, fmt::Debug};
 
@@ -278,6 +279,15 @@ impl<'b, TX: DbTx> StateRootProvider for HistoricalStateProviderRef<'b, TX> {
         StateRoot::overlay_root_with_updates(self.tx, revert_state, Default::default())
             .map_err(|err| ProviderError::Database(err.into()))
     }
+
+    fn hashed_storage_root(
+        &self,
+        address: Address,
+        hashed_storage: HashedStorage,
+    ) -> ProviderResult<B256> {
+        StorageRoot::overlay_root(self.tx, address, hashed_storage)
+            .map_err(|err| ProviderError::Database(err.into()))
+    }
 }
 
 impl<'b, TX: DbTx> StateProofProvider for HistoricalStateProviderRef<'b, TX> {

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

@@ -17,9 +17,9 @@ use reth_storage_api::StateProofProvider;
 use reth_storage_errors::provider::{ProviderError, ProviderResult};
 use reth_trie::{
     proof::Proof, updates::TrieUpdates, witness::TrieWitness, AccountProof, HashedPostState,
-    StateRoot,
+    HashedStorage, StateRoot, StorageRoot,
 };
-use reth_trie_db::{DatabaseProof, DatabaseStateRoot, DatabaseTrieWitness};
+use reth_trie_db::{DatabaseProof, DatabaseStateRoot, DatabaseStorageRoot, DatabaseTrieWitness};
 
 /// State provider over latest state that takes tx reference.
 #[derive(Debug)]
@@ -93,6 +93,15 @@ impl<'b, TX: DbTx> StateRootProvider for LatestStateProviderRef<'b, TX> {
         StateRoot::overlay_root_with_updates(self.tx, hashed_state, Default::default())
             .map_err(|err| ProviderError::Database(err.into()))
     }
+
+    fn hashed_storage_root(
+        &self,
+        address: Address,
+        hashed_storage: HashedStorage,
+    ) -> ProviderResult<B256> {
+        StorageRoot::overlay_root(self.tx, address, hashed_storage)
+            .map_err(|err| ProviderError::Database(err.into()))
+    }
 }
 
 impl<'b, TX: DbTx> StateProofProvider for LatestStateProviderRef<'b, TX> {

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

@@ -46,6 +46,7 @@ macro_rules! delegate_provider_impls {
                 fn hashed_state_root(&self, state: reth_trie::HashedPostState) -> reth_storage_errors::provider::ProviderResult<reth_primitives::B256>;
                 fn state_root_with_updates(&self, state: &revm::db::BundleState) -> reth_storage_errors::provider::ProviderResult<(reth_primitives::B256, reth_trie::updates::TrieUpdates)>;
                 fn hashed_state_root_with_updates(&self, state: reth_trie::HashedPostState) -> reth_storage_errors::provider::ProviderResult<(reth_primitives::B256, reth_trie::updates::TrieUpdates)>;
+                fn hashed_storage_root(&self, address: reth_primitives::Address, storage: reth_trie::HashedStorage) ->  reth_storage_errors::provider::ProviderResult<reth_primitives::B256>;
             }
             StateProofProvider $(where [$($generics)*])? {
                 fn proof(&self, state: &revm::db::BundleState, address: reth_primitives::Address, slots: &[reth_primitives::B256]) -> reth_storage_errors::provider::ProviderResult<reth_trie::AccountProof>;

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

@@ -17,7 +17,7 @@ use reth_primitives::{
 };
 use reth_storage_api::StateProofProvider;
 use reth_storage_errors::provider::{ProviderError, ProviderResult};
-use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState};
+use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage};
 use revm::primitives::{BlockEnv, CfgEnvWithHandlerCfg};
 use std::{
     collections::{BTreeMap, HashMap},
@@ -559,6 +559,14 @@ impl StateRootProvider for MockEthProvider {
         let state_root = self.state_roots.lock().pop().unwrap_or_default();
         Ok((state_root, Default::default()))
     }
+
+    fn hashed_storage_root(
+        &self,
+        _address: Address,
+        _hashed_storage: HashedStorage,
+    ) -> ProviderResult<B256> {
+        Ok(B256::default())
+    }
 }
 
 impl StateProofProvider for MockEthProvider {

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

@@ -323,6 +323,14 @@ impl StateRootProvider for NoopProvider {
     ) -> ProviderResult<(B256, TrieUpdates)> {
         Ok((B256::default(), TrieUpdates::default()))
     }
+
+    fn hashed_storage_root(
+        &self,
+        _address: Address,
+        _hashed_storage: reth_trie::HashedStorage,
+    ) -> ProviderResult<B256> {
+        Ok(B256::default())
+    }
 }
 
 impl StateProofProvider for NoopProvider {

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

@@ -1,6 +1,6 @@
 use reth_primitives::{Address, Bytes, B256};
 use reth_storage_errors::provider::ProviderResult;
-use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState};
+use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage};
 use revm::db::BundleState;
 use std::collections::HashMap;
 
@@ -36,6 +36,14 @@ pub trait StateRootProvider: Send + Sync {
         &self,
         hashed_state: HashedPostState,
     ) -> ProviderResult<(B256, TrieUpdates)>;
+
+    /// Returns the storage root of the `HashedStorage` for target address on top of the current
+    /// state.
+    fn hashed_storage_root(
+        &self,
+        address: Address,
+        hashed_storage: HashedStorage,
+    ) -> ProviderResult<B256>;
 }
 
 /// A type that can generate state proof on top of a given post state.

crates/trie/db/src/lib.rs

@@ -14,7 +14,7 @@ pub use hashed_cursor::{
 pub use prefix_set::PrefixSetLoader;
 pub use proof::DatabaseProof;
 pub use state::{DatabaseHashedPostState, DatabaseStateRoot};
-pub use storage::DatabaseStorageRoot;
+pub use storage::{DatabaseHashedStorage, DatabaseStorageRoot};
 pub use trie_cursor::{
     DatabaseAccountTrieCursor, DatabaseStorageTrieCursor, DatabaseTrieCursorFactory,
 };

crates/trie/db/src/storage.rs

@@ -1,7 +1,11 @@
 use crate::{DatabaseHashedCursorFactory, DatabaseTrieCursorFactory};
+use reth_db::{cursor::DbCursorRO, models::BlockNumberAddress, tables, DatabaseError};
 use reth_db_api::transaction::DbTx;
-use reth_primitives::{Address, B256};
+use reth_execution_errors::StorageRootError;
-use reth_trie::StorageRoot;
+use reth_primitives::{keccak256, Address, BlockNumber, B256};
+use reth_trie::{
+    hashed_cursor::HashedPostStateCursorFactory, HashedPostState, HashedStorage, StorageRoot,
+};
 
 #[cfg(feature = "metrics")]
 use reth_trie::metrics::{TrieRootMetrics, TrieType};
@@ -13,6 +17,20 @@ pub trait DatabaseStorageRoot<'a, TX> {
 
     /// Create a new storage root calculator from database transaction and hashed address.
     fn from_tx_hashed(tx: &'a TX, hashed_address: B256) -> Self;
+
+    /// Calculates the storage root for this [`HashedStorage`] and returns it.
+    fn overlay_root(
+        tx: &'a TX,
+        address: Address,
+        hashed_storage: HashedStorage,
+    ) -> Result<B256, StorageRootError>;
+}
+
+/// Extends [`HashedStorage`] with operations specific for working with a database transaction.
+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>;
 }
 
 impl<'a, TX: DbTx> DatabaseStorageRoot<'a, TX>
@@ -37,4 +55,38 @@ impl<'a, TX: DbTx> DatabaseStorageRoot<'a, TX>
             TrieRootMetrics::new(TrieType::Storage),
         )
     }
+
+    fn overlay_root(
+        tx: &'a TX,
+        address: Address,
+        hashed_storage: HashedStorage,
+    ) -> Result<B256, StorageRootError> {
+        let prefix_set = hashed_storage.construct_prefix_set().freeze();
+        let state_sorted =
+            HashedPostState::from_hashed_storage(keccak256(address), hashed_storage).into_sorted();
+        StorageRoot::new(
+            DatabaseTrieCursorFactory::new(tx),
+            HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &state_sorted),
+            address,
+            #[cfg(feature = "metrics")]
+            TrieRootMetrics::new(TrieType::Storage),
+        )
+        .with_prefix_set(prefix_set)
+        .root()
+    }
+}
+
+impl<TX: DbTx> DatabaseHashedStorage<TX> for HashedStorage {
+    fn from_reverts(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);
+                storage.storage.insert(hashed_slot, storage_change.value);
+            }
+        }
+        Ok(storage)
+    }
 }

crates/trie/trie/src/state.rs

@@ -5,7 +5,7 @@ use crate::{
 use itertools::Itertools;
 use rayon::prelude::{IntoParallelIterator, ParallelIterator};
 use reth_primitives::{keccak256, Account, Address, B256, U256};
-use revm::db::BundleAccount;
+use revm::db::{states::StorageSlot, AccountStatus, BundleAccount};
 use std::collections::{hash_map, HashMap, HashSet};
 
 /// Representation of in-memory hashed state.
@@ -29,12 +29,8 @@ impl HashedPostState {
             .map(|(address, account)| {
                 let hashed_address = keccak256(address);
                 let hashed_account = account.info.clone().map(Into::into);
-                let hashed_storage = HashedStorage::from_iter(
+                let hashed_storage =
-                    account.status.was_destroyed(),
+                    HashedStorage::from_bundle_state(account.status, &account.storage);
-                    account.storage.iter().map(|(key, value)| {
-                        (keccak256(B256::new(key.to_be_bytes())), value.present_value)
-                    }),
-                );
                 (hashed_address, (hashed_account, hashed_storage))
             })
             .collect::<Vec<(B256, (Option<Account>, HashedStorage))>>();
@@ -48,6 +44,11 @@ impl HashedPostState {
         Self { accounts, storages }
     }
 
+    /// Construct [`HashedPostState`] from a single [`HashedStorage`].
+    pub fn from_hashed_storage(hashed_address: B256, storage: HashedStorage) -> Self {
+        Self { accounts: HashMap::default(), storages: HashMap::from([(hashed_address, storage)]) }
+    }
+
     /// Set account entries on hashed state.
     pub fn with_accounts(
         mut self,
@@ -127,12 +128,7 @@ impl HashedPostState {
         let mut storage_prefix_sets = HashMap::with_capacity(self.storages.len());
         for (hashed_address, hashed_storage) in &self.storages {
             account_prefix_set.insert(Nibbles::unpack(hashed_address));
-
+            storage_prefix_sets.insert(*hashed_address, hashed_storage.construct_prefix_set());
-            let mut prefix_set = PrefixSetMut::with_capacity(hashed_storage.storage.len());
-            for hashed_slot in hashed_storage.storage.keys() {
-                prefix_set.insert(Nibbles::unpack(hashed_slot));
-            }
-            storage_prefix_sets.insert(*hashed_address, prefix_set);
         }
 
         TriePrefixSetsMut { account_prefix_set, storage_prefix_sets, destroyed_accounts }
@@ -159,6 +155,24 @@ impl HashedStorage {
         Self { wiped, storage: HashMap::from_iter(iter) }
     }
 
+    /// Create new hashed storage from bundle state account entry.
+    pub fn from_bundle_state(status: AccountStatus, storage: &HashMap<U256, StorageSlot>) -> Self {
+        let storage = storage
+            .iter()
+            .map(|(key, value)| (keccak256(B256::from(*key)), value.present_value))
+            .collect();
+        Self { wiped: status.was_destroyed(), storage }
+    }
+
+    /// Construct [`PrefixSetMut`] from hashed storage.
+    pub fn construct_prefix_set(&self) -> PrefixSetMut {
+        let mut prefix_set = PrefixSetMut::with_capacity(self.storage.len());
+        for hashed_slot in self.storage.keys() {
+            prefix_set.insert(Nibbles::unpack(hashed_slot));
+        }
+        prefix_set
+    }
+
     /// Extend hashed storage with contents of other.
     /// The entries in second hashed storage take precedence.
     pub fn extend(&mut self, other: Self) {
@@ -198,6 +212,14 @@ pub struct HashedPostStateSorted {
 }
 
 impl HashedPostStateSorted {
+    /// Create new instance of [`HashedPostStateSorted`]
+    pub const fn new(
+        accounts: HashedAccountsSorted,
+        storages: HashMap<B256, HashedStorageSorted>,
+    ) -> Self {
+        Self { accounts, storages }
+    }
+
     /// Returns reference to hashed accounts.
     pub const fn accounts(&self) -> &HashedAccountsSorted {
         &self.accounts