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

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

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_trie::StorageRoot;
+use reth_execution_errors::StorageRootError;
+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(
-                    account.status.was_destroyed(),
-                    account.storage.iter().map(|(key, value)| {
-                        (keccak256(B256::new(key.to_be_bytes())), value.present_value)
-                    }),
-                );
+                let hashed_storage =
+                    HashedStorage::from_bundle_state(account.status, &account.storage);
                 (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));
-
-            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);
+            storage_prefix_sets.insert(*hashed_address, hashed_storage.construct_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