chore(trie): account specific hashed storage cursor (#8377)

crates/trie/src/hashed_cursor/default.rs

@@ -8,14 +8,21 @@ use reth_primitives::{Account, StorageEntry, B256};
 
 impl<'a, TX: DbTx> HashedCursorFactory for &'a TX {
     type AccountCursor = <TX as DbTx>::Cursor<tables::HashedAccounts>;
-    type StorageCursor = <TX as DbTx>::DupCursor<tables::HashedStorages>;
+    type StorageCursor =
+        DatabaseHashedStorageCursor<<TX as DbTx>::DupCursor<tables::HashedStorages>>;
 
     fn hashed_account_cursor(&self) -> Result<Self::AccountCursor, reth_db::DatabaseError> {
         self.cursor_read::<tables::HashedAccounts>()
     }
 
-    fn hashed_storage_cursor(&self) -> Result<Self::StorageCursor, reth_db::DatabaseError> {
-        self.cursor_dup_read::<tables::HashedStorages>()
+    fn hashed_storage_cursor(
+        &self,
+        hashed_address: B256,
+    ) -> Result<Self::StorageCursor, reth_db::DatabaseError> {
+        Ok(DatabaseHashedStorageCursor::new(
+            self.cursor_dup_read::<tables::HashedStorages>()?,
+            hashed_address,
+        ))
     }
 }
 
@@ -32,23 +39,35 @@ where
     }
 }
 
-impl<C> HashedStorageCursor for C
+/// The structure wrapping a database cursor for hashed storage and
+/// a target hashed address. Implements [HashedStorageCursor] for iterating
+/// hashed state
+#[derive(Debug)]
+pub struct DatabaseHashedStorageCursor<C> {
+    cursor: C,
+    hashed_address: B256,
+}
+
+impl<C> DatabaseHashedStorageCursor<C> {
+    /// Create new [DatabaseHashedStorageCursor].
+    pub fn new(cursor: C, hashed_address: B256) -> Self {
+        Self { cursor, hashed_address }
+    }
+}
+
+impl<C> HashedStorageCursor for DatabaseHashedStorageCursor<C>
 where
     C: DbCursorRO<tables::HashedStorages> + DbDupCursorRO<tables::HashedStorages>,
 {
-    fn is_storage_empty(&mut self, key: B256) -> Result<bool, reth_db::DatabaseError> {
-        Ok(self.seek_exact(key)?.is_none())
+    fn is_storage_empty(&mut self) -> Result<bool, reth_db::DatabaseError> {
+        Ok(self.cursor.seek_exact(self.hashed_address)?.is_none())
     }
 
-    fn seek(
-        &mut self,
-        key: B256,
-        subkey: B256,
-    ) -> Result<Option<StorageEntry>, reth_db::DatabaseError> {
-        self.seek_by_key_subkey(key, subkey)
+    fn seek(&mut self, subkey: B256) -> Result<Option<StorageEntry>, reth_db::DatabaseError> {
+        self.cursor.seek_by_key_subkey(self.hashed_address, subkey)
     }
 
     fn next(&mut self) -> Result<Option<StorageEntry>, reth_db::DatabaseError> {
-        self.next_dup_val()
+        self.cursor.next_dup_val()
     }
 }

crates/trie/src/hashed_cursor/mod.rs

@@ -2,6 +2,7 @@ use reth_primitives::{Account, StorageEntry, B256};
 
 /// Default implementation of the hashed state cursor traits.
 mod default;
+pub use default::DatabaseHashedStorageCursor;
 
 /// Implementation of hashed state cursor traits for the post state.
 mod post_state;
@@ -18,7 +19,10 @@ pub trait HashedCursorFactory {
     fn hashed_account_cursor(&self) -> Result<Self::AccountCursor, reth_db::DatabaseError>;
 
     /// Returns a cursor for iterating over all hashed storage entries in the state.
-    fn hashed_storage_cursor(&self) -> Result<Self::StorageCursor, reth_db::DatabaseError>;
+    fn hashed_storage_cursor(
+        &self,
+        hashed_address: B256,
+    ) -> Result<Self::StorageCursor, reth_db::DatabaseError>;
 }
 
 /// The cursor for iterating over hashed accounts.
@@ -33,14 +37,10 @@ pub trait HashedAccountCursor {
 /// The cursor for iterating over hashed storage entries.
 pub trait HashedStorageCursor {
     /// Returns `true` if there are no entries for a given key.
-    fn is_storage_empty(&mut self, key: B256) -> Result<bool, reth_db::DatabaseError>;
+    fn is_storage_empty(&mut self) -> Result<bool, reth_db::DatabaseError>;
 
     /// Seek an entry greater or equal to the given key/subkey and position the cursor there.
-    fn seek(
-        &mut self,
-        key: B256,
-        subkey: B256,
-    ) -> Result<Option<StorageEntry>, reth_db::DatabaseError>;
+    fn seek(&mut self, subkey: B256) -> Result<Option<StorageEntry>, reth_db::DatabaseError>;
 
     /// Move the cursor to the next entry and return it.
     fn next(&mut self) -> Result<Option<StorageEntry>, reth_db::DatabaseError>;

crates/trie/src/hashed_cursor/post_state.rs

@@ -25,9 +25,12 @@ impl<'a, CF: HashedCursorFactory> HashedCursorFactory for HashedPostStateCursorF
         Ok(HashedPostStateAccountCursor::new(cursor, self.post_state))
     }
 
-    fn hashed_storage_cursor(&self) -> Result<Self::StorageCursor, reth_db::DatabaseError> {
-        let cursor = self.cursor_factory.hashed_storage_cursor()?;
-        Ok(HashedPostStateStorageCursor::new(cursor, self.post_state))
+    fn hashed_storage_cursor(
+        &self,
+        hashed_address: B256,
+    ) -> Result<Self::StorageCursor, reth_db::DatabaseError> {
+        let cursor = self.cursor_factory.hashed_storage_cursor(hashed_address)?;
+        Ok(HashedPostStateStorageCursor::new(cursor, self.post_state, hashed_address))
     }
 }
 
@@ -179,10 +182,10 @@ pub struct HashedPostStateStorageCursor<'b, C> {
     cursor: C,
     /// The reference to the post state.
     post_state: &'b HashedPostStateSorted,
+    /// The current hashed account key.
+    hashed_address: B256,
     /// The post state index where the cursor is currently at.
     post_state_storage_index: usize,
-    /// The current hashed account key.
-    account: Option<B256>,
     /// The last slot that has been returned by the cursor.
     /// De facto, this is the cursor's position for the given account key.
     last_slot: Option<B256>,
@@ -190,14 +193,14 @@ pub struct HashedPostStateStorageCursor<'b, C> {
 
 impl<'b, C> HashedPostStateStorageCursor<'b, C> {
     /// Create new instance of [HashedPostStateStorageCursor].
-    pub fn new(cursor: C, post_state: &'b HashedPostStateSorted) -> Self {
-        Self { cursor, post_state, account: None, last_slot: None, post_state_storage_index: 0 }
+    pub fn new(cursor: C, post_state: &'b HashedPostStateSorted, hashed_address: B256) -> Self {
+        Self { cursor, post_state, hashed_address, last_slot: None, post_state_storage_index: 0 }
     }
 
     /// Returns `true` if the storage for the given
     /// The database is not checked since it already has no wiped storage entries.
-    fn is_db_storage_wiped(&self, account: &B256) -> bool {
-        match self.post_state.storages.get(account) {
+    fn is_db_storage_wiped(&self) -> bool {
+        match self.post_state.storages.get(&self.hashed_address) {
             Some(storage) => storage.wiped,
             None => false,
         }
@@ -205,10 +208,10 @@ impl<'b, C> HashedPostStateStorageCursor<'b, C> {
 
     /// Check if the slot was zeroed out in the post state.
     /// The database is not checked since it already has no zero-valued slots.
-    fn is_slot_zero_valued(&self, account: &B256, slot: &B256) -> bool {
+    fn is_slot_zero_valued(&self, slot: &B256) -> bool {
         self.post_state
             .storages
-            .get(account)
+            .get(&self.hashed_address)
             .map(|storage| storage.zero_valued_slots.contains(slot))
             .unwrap_or_default()
     }
@@ -247,34 +250,24 @@ where
     ///
     /// This function should be called before attempting to call [HashedStorageCursor::seek] or
     /// [HashedStorageCursor::next].
-    fn is_storage_empty(&mut self, key: B256) -> Result<bool, reth_db::DatabaseError> {
-        let is_empty = match self.post_state.storages.get(&key) {
+    fn is_storage_empty(&mut self) -> Result<bool, reth_db::DatabaseError> {
+        let is_empty = match self.post_state.storages.get(&self.hashed_address) {
             Some(storage) => {
                 // If the storage has been wiped at any point
                 storage.wiped &&
                     // and the current storage does not contain any non-zero values
                     storage.non_zero_valued_slots.is_empty()
             }
-            None => self.cursor.is_storage_empty(key)?,
+            None => self.cursor.is_storage_empty()?,
         };
         Ok(is_empty)
     }
 
     /// Seek the next account storage entry for a given hashed key pair.
-    fn seek(
-        &mut self,
-        account: B256,
-        subkey: B256,
-    ) -> Result<Option<StorageEntry>, reth_db::DatabaseError> {
-        if self.account.map_or(true, |acc| acc != account) {
-            self.account = Some(account);
-            self.last_slot = None;
-            self.post_state_storage_index = 0;
-        }
-
+    fn seek(&mut self, subkey: B256) -> Result<Option<StorageEntry>, reth_db::DatabaseError> {
         // Attempt to find the account's storage in post state.
         let mut post_state_entry = None;
-        if let Some(storage) = self.post_state.storages.get(&account) {
+        if let Some(storage) = self.post_state.storages.get(&self.hashed_address) {
             post_state_entry = storage.non_zero_valued_slots.get(self.post_state_storage_index);
 
             while post_state_entry.map(|(slot, _)| slot < &subkey).unwrap_or_default() {
@@ -293,14 +286,14 @@ where
         }
 
         // It's not an exact match, reposition to the first greater or equal account.
-        let db_entry = if self.is_db_storage_wiped(&account) {
+        let db_entry = if self.is_db_storage_wiped() {
             None
         } else {
-            let mut db_entry = self.cursor.seek(account, subkey)?;
+            let mut db_entry = self.cursor.seek(subkey)?;
 
             while db_entry
                 .as_ref()
-                .map(|entry| self.is_slot_zero_valued(&account, &entry.key))
+                .map(|entry| self.is_slot_zero_valued(&entry.key))
                 .unwrap_or_default()
             {
                 db_entry = self.cursor.next()?;
@@ -322,25 +315,21 @@ where
     /// If the account key is not set. [HashedStorageCursor::seek] must be called first in order to
     /// position the cursor.
     fn next(&mut self) -> Result<Option<StorageEntry>, reth_db::DatabaseError> {
-        let account = self.account.expect("`seek` must be called first");
-
         let last_slot = match self.last_slot.as_ref() {
             Some(slot) => slot,
             None => return Ok(None), // no previous entry was found
         };
 
-        let db_entry = if self.is_db_storage_wiped(&account) {
+        let db_entry = if self.is_db_storage_wiped() {
             None
         } else {
             // If post state was given precedence, move the cursor forward.
-            let mut db_entry = self.cursor.seek(account, *last_slot)?;
+            let mut db_entry = self.cursor.seek(*last_slot)?;
 
             // If the entry was already returned or is zero-values, move to the next.
             while db_entry
                 .as_ref()
-                .map(|entry| {
-                    &entry.key == last_slot || self.is_slot_zero_valued(&account, &entry.key)
-                })
+                .map(|entry| &entry.key == last_slot || self.is_slot_zero_valued(&entry.key))
                 .unwrap_or_default()
             {
                 db_entry = self.cursor.next()?;
@@ -351,7 +340,7 @@ where
 
         // Attempt to find the account's storage in post state.
         let mut post_state_entry = None;
-        if let Some(storage) = self.post_state.storages.get(&account) {
+        if let Some(storage) = self.post_state.storages.get(&self.hashed_address) {
             post_state_entry = storage.non_zero_valued_slots.get(self.post_state_storage_index);
             while post_state_entry.map(|(slot, _)| slot <= last_slot).unwrap_or_default() {
                 self.post_state_storage_index += 1;
@@ -397,12 +386,11 @@ mod tests {
         factory: &impl HashedCursorFactory,
         expected: impl Iterator<Item = (B256, BTreeMap<B256, U256>)>,
     ) {
-        let mut cursor = factory.hashed_storage_cursor().unwrap();
-
         for (account, storage) in expected {
+            let mut cursor = factory.hashed_storage_cursor(account).unwrap();
             let mut expected_storage = storage.into_iter();
 
-            let first_storage = cursor.seek(account, B256::default()).unwrap();
+            let first_storage = cursor.seek(B256::default()).unwrap();
             assert_eq!(first_storage.map(|e| (e.key, e.value)), expected_storage.next());
 
             for expected_entry in expected_storage {
@@ -577,8 +565,8 @@ mod tests {
             let sorted = HashedPostState::default().into_sorted();
             let tx = db.tx().unwrap();
             let factory = HashedPostStateCursorFactory::new(&tx, &sorted);
-            let mut cursor = factory.hashed_storage_cursor().unwrap();
-            assert!(cursor.is_storage_empty(address).unwrap());
+            let mut cursor = factory.hashed_storage_cursor(address).unwrap();
+            assert!(cursor.is_storage_empty().unwrap());
         }
 
         let db_storage =
@@ -600,8 +588,8 @@ mod tests {
             let sorted = HashedPostState::default().into_sorted();
             let tx = db.tx().unwrap();
             let factory = HashedPostStateCursorFactory::new(&tx, &sorted);
-            let mut cursor = factory.hashed_storage_cursor().unwrap();
-            assert!(!cursor.is_storage_empty(address).unwrap());
+            let mut cursor = factory.hashed_storage_cursor(address).unwrap();
+            assert!(!cursor.is_storage_empty().unwrap());
         }
 
         // wiped storage, must be empty
@@ -615,8 +603,8 @@ mod tests {
             let sorted = hashed_post_state.into_sorted();
             let tx = db.tx().unwrap();
             let factory = HashedPostStateCursorFactory::new(&tx, &sorted);
-            let mut cursor = factory.hashed_storage_cursor().unwrap();
-            assert!(cursor.is_storage_empty(address).unwrap());
+            let mut cursor = factory.hashed_storage_cursor(address).unwrap();
+            assert!(cursor.is_storage_empty().unwrap());
         }
 
         // wiped storage, but post state has zero-value entries
@@ -631,8 +619,8 @@ mod tests {
             let sorted = hashed_post_state.into_sorted();
             let tx = db.tx().unwrap();
             let factory = HashedPostStateCursorFactory::new(&tx, &sorted);
-            let mut cursor = factory.hashed_storage_cursor().unwrap();
-            assert!(cursor.is_storage_empty(address).unwrap());
+            let mut cursor = factory.hashed_storage_cursor(address).unwrap();
+            assert!(cursor.is_storage_empty().unwrap());
         }
 
         // wiped storage, but post state has non-zero entries
@@ -647,8 +635,8 @@ mod tests {
             let sorted = hashed_post_state.into_sorted();
             let tx = db.tx().unwrap();
             let factory = HashedPostStateCursorFactory::new(&tx, &sorted);
-            let mut cursor = factory.hashed_storage_cursor().unwrap();
-            assert!(!cursor.is_storage_empty(address).unwrap());
+            let mut cursor = factory.hashed_storage_cursor(address).unwrap();
+            assert!(!cursor.is_storage_empty().unwrap());
         }
     }
 

crates/trie/src/node_iter.rs

@@ -166,8 +166,6 @@ pub struct StorageNodeIter<C, H> {
     pub walker: TrieWalker<C>,
     /// The cursor for the hashed storage entries.
     pub hashed_storage_cursor: H,
-    /// The hashed address this storage trie belongs to.
-    hashed_address: B256,
 
     /// Current hashed storage entry.
     current_hashed_entry: Option<StorageEntry>,
@@ -177,11 +175,10 @@ pub struct StorageNodeIter<C, H> {
 
 impl<C, H> StorageNodeIter<C, H> {
     /// Creates a new instance of StorageNodeIter.
-    pub fn new(walker: TrieWalker<C>, hashed_storage_cursor: H, hashed_address: B256) -> Self {
+    pub fn new(walker: TrieWalker<C>, hashed_storage_cursor: H) -> Self {
         Self {
             walker,
             hashed_storage_cursor,
-            hashed_address,
             current_walker_key_checked: false,
             current_hashed_entry: None,
         }
@@ -238,8 +235,7 @@ where
             // Attempt to get the next unprocessed key from the walker.
             if let Some(seek_key) = self.walker.next_unprocessed_key() {
                 // Seek and update the current hashed entry based on the new seek key.
-                self.current_hashed_entry =
-                    self.hashed_storage_cursor.seek(self.hashed_address, seek_key)?;
+                self.current_hashed_entry = self.hashed_storage_cursor.seek(seek_key)?;
                 self.walker.advance()?;
             } else {
                 // No more keys to process, break the loop.

crates/trie/src/proof.rs

@@ -109,12 +109,13 @@ where
         hashed_address: B256,
         slots: &[B256],
     ) -> Result<(B256, Vec<StorageProof>), StorageRootError> {
-        let mut hashed_storage_cursor = self.hashed_cursor_factory.hashed_storage_cursor()?;
+        let mut hashed_storage_cursor =
+            self.hashed_cursor_factory.hashed_storage_cursor(hashed_address)?;
 
         let mut proofs = slots.iter().copied().map(StorageProof::new).collect::<Vec<_>>();
 
         // short circuit on empty storage
-        if hashed_storage_cursor.is_storage_empty(hashed_address)? {
+        if hashed_storage_cursor.is_storage_empty()? {
             return Ok((EMPTY_ROOT_HASH, proofs))
         }
 
@@ -128,8 +129,7 @@ where
 
         let retainer = ProofRetainer::from_iter(target_nibbles);
         let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer);
-        let mut storage_node_iter =
-            StorageNodeIter::new(walker, hashed_storage_cursor, hashed_address);
+        let mut storage_node_iter = StorageNodeIter::new(walker, hashed_storage_cursor);
         while let Some(node) = storage_node_iter.try_next()? {
             match node {
                 StorageNode::Branch(node) => {

crates/trie/src/trie.rs

@@ -483,10 +483,11 @@ where
     ) -> Result<(B256, usize, TrieUpdates), StorageRootError> {
         trace!(target: "trie::storage_root", hashed_address = ?self.hashed_address, "calculating storage root");
 
-        let mut hashed_storage_cursor = self.hashed_cursor_factory.hashed_storage_cursor()?;
+        let mut hashed_storage_cursor =
+            self.hashed_cursor_factory.hashed_storage_cursor(self.hashed_address)?;
 
         // short circuit on empty storage
-        if hashed_storage_cursor.is_storage_empty(self.hashed_address)? {
+        if hashed_storage_cursor.is_storage_empty()? {
             return Ok((
                 EMPTY_ROOT_HASH,
                 0,
@@ -500,8 +501,7 @@ where
 
         let mut hash_builder = HashBuilder::default().with_updates(retain_updates);
 
-        let mut storage_node_iter =
-            StorageNodeIter::new(walker, hashed_storage_cursor, self.hashed_address);
+        let mut storage_node_iter = StorageNodeIter::new(walker, hashed_storage_cursor);
         while let Some(node) = storage_node_iter.try_next()? {
             match node {
                 StorageNode::Branch(node) => {