refactor(trie): trie node iterators (#5048)

Cargo.lock

@@ -6448,6 +6448,7 @@ name = "reth-trie"
 version = "0.1.0-alpha.10"
 dependencies = [
  "alloy-rlp",
+ "auto_impl",
  "criterion",
  "derive_more",
  "once_cell",

crates/primitives/src/stage/checkpoints.rs

@@ -17,6 +17,7 @@ pub struct MerkleCheckpoint {
     pub target_block: BlockNumber,
     /// The last hashed account key processed.
     pub last_account_key: B256,
+    // TODO: remove in the next breaking release.
     /// The last walker key processed.
     pub last_walker_key: Vec<u8>,
     /// Previously recorded walker stack.
@@ -30,11 +31,16 @@ impl MerkleCheckpoint {
     pub fn new(
         target_block: BlockNumber,
         last_account_key: B256,
-        last_walker_key: Vec<u8>,
         walker_stack: Vec<StoredSubNode>,
         state: HashBuilderState,
     ) -> Self {
-        Self { target_block, last_account_key, last_walker_key, walker_stack, state }
+        Self {
+            target_block,
+            last_account_key,
+            walker_stack,
+            state,
+            last_walker_key: Vec::default(),
+        }
     }
 }
 

crates/stages/src/stages/merkle.rs

@@ -224,7 +224,6 @@ impl<DB: Database> Stage<DB> for MerkleStage {
                     let checkpoint = MerkleCheckpoint::new(
                         to_block,
                         state.last_account_key,
-                        state.last_walker_key.hex_data.to_vec(),
                         state.walker_stack.into_iter().map(StoredSubNode::from).collect(),
                         state.hash_builder.into(),
                     );

crates/trie/Cargo.toml

@@ -27,6 +27,7 @@ tracing.workspace = true
 # misc 
 thiserror.workspace = true
 derive_more = "0.99"
+auto_impl = "1"
 
 # test-utils
 triehash = { version = "0.8", optional = true }

crates/trie/src/lib.rs

@@ -34,6 +34,9 @@ pub mod walker;
 mod errors;
 pub use errors::*;
 
+// The iterators for traversing existing intermediate hashes and updated trie leaves.
+pub(crate) mod node_iter;
+
 /// Merkle proof generation.
 pub mod proof;
 

crates/trie/src/node_iter.rs

@@ -0,0 +1,208 @@
+use crate::{
+    hashed_cursor::{HashedAccountCursor, HashedStorageCursor},
+    trie_cursor::TrieCursor,
+    walker::TrieWalker,
+    StateRootError, StorageRootError,
+};
+use reth_primitives::{trie::Nibbles, Account, StorageEntry, B256, U256};
+
+#[derive(Debug)]
+pub(crate) struct TrieBranchNode {
+    pub(crate) key: Nibbles,
+    pub(crate) value: B256,
+    pub(crate) children_are_in_trie: bool,
+}
+
+impl TrieBranchNode {
+    pub(crate) fn new(key: Nibbles, value: B256, children_are_in_trie: bool) -> Self {
+        Self { key, value, children_are_in_trie }
+    }
+}
+
+#[derive(Debug)]
+pub(crate) enum AccountNode {
+    Branch(TrieBranchNode),
+    Leaf(B256, Account),
+}
+
+#[derive(Debug)]
+pub(crate) enum StorageNode {
+    Branch(TrieBranchNode),
+    Leaf(B256, U256),
+}
+
+/// An iterator over existing intermediate branch nodes and updated leaf nodes.
+#[derive(Debug)]
+pub(crate) struct AccountNodeIter<C, H> {
+    /// Underlying walker over intermediate nodes.
+    pub(crate) walker: TrieWalker<C>,
+    /// The cursor for the hashed account entries.
+    pub(crate) hashed_account_cursor: H,
+    /// The previous account key. If the iteration was previously interrupted, this value can be
+    /// used to resume iterating from the last returned leaf node.
+    previous_account_key: Option<B256>,
+
+    /// Current hashed account entry.
+    current_hashed_entry: Option<(B256, Account)>,
+    /// Flag indicating whether we should check the current walker key.
+    current_walker_key_checked: bool,
+}
+
+impl<C, H> AccountNodeIter<C, H> {
+    pub(crate) fn new(walker: TrieWalker<C>, hashed_account_cursor: H) -> Self {
+        Self {
+            walker,
+            hashed_account_cursor,
+            previous_account_key: None,
+            current_hashed_entry: None,
+            current_walker_key_checked: false,
+        }
+    }
+
+    pub(crate) fn with_last_account_key(mut self, previous_account_key: B256) -> Self {
+        self.previous_account_key = Some(previous_account_key);
+        self
+    }
+}
+
+impl<C, H> AccountNodeIter<C, H>
+where
+    C: TrieCursor,
+    H: HashedAccountCursor,
+{
+    /// Return the next account trie node to be added to the hash builder.
+    ///
+    /// Returns the nodes using this algorithm:
+    /// 1. Return the current intermediate branch node if it hasn't been updated.
+    /// 2. Advance the trie walker to the next intermediate branch node and retrieve next
+    ///    unprocessed key.
+    /// 3. Reposition the hashed account cursor on the next unprocessed key.
+    /// 4. Return every hashed account entry up to the key of the current intermediate branch node.
+    /// 5. Repeat.
+    ///
+    /// NOTE: The iteration will start from the key of the previous hashed entry if it was supplied.
+    pub(crate) fn try_next(&mut self) -> Result<Option<AccountNode>, StateRootError> {
+        loop {
+            if let Some(key) = self.walker.key() {
+                if !self.current_walker_key_checked && self.previous_account_key.is_none() {
+                    self.current_walker_key_checked = true;
+                    if self.walker.can_skip_current_node {
+                        return Ok(Some(AccountNode::Branch(TrieBranchNode::new(
+                            key,
+                            self.walker.hash().unwrap(),
+                            self.walker.children_are_in_trie(),
+                        ))))
+                    }
+                }
+            }
+
+            if let Some((hashed_address, account)) = self.current_hashed_entry.take() {
+                if self.walker.key().map_or(false, |key| key < Nibbles::unpack(hashed_address)) {
+                    self.current_walker_key_checked = false;
+                    continue
+                }
+
+                self.current_hashed_entry = self.hashed_account_cursor.next()?;
+                return Ok(Some(AccountNode::Leaf(hashed_address, account)))
+            }
+
+            match self.previous_account_key.take() {
+                Some(account_key) => {
+                    self.hashed_account_cursor.seek(account_key)?;
+                    self.current_hashed_entry = self.hashed_account_cursor.next()?;
+                }
+                None => {
+                    let seek_key = match self.walker.next_unprocessed_key() {
+                        Some(key) => key,
+                        None => break, // no more keys
+                    };
+                    self.current_hashed_entry = self.hashed_account_cursor.seek(seek_key)?;
+                    self.walker.advance()?;
+                }
+            }
+        }
+
+        Ok(None)
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct StorageNodeIter<C, H> {
+    /// Underlying walker over intermediate nodes.
+    pub(crate) walker: TrieWalker<C>,
+    /// The cursor for the hashed storage entries.
+    pub(crate) hashed_storage_cursor: H,
+    /// The hashed address this storage trie belongs to.
+    hashed_address: B256,
+
+    /// Current hashed storage entry.
+    current_hashed_entry: Option<StorageEntry>,
+    /// Flag indicating whether we should check the current walker key.
+    current_walker_key_checked: bool,
+}
+
+impl<C, H> StorageNodeIter<C, H> {
+    pub(crate) fn new(
+        walker: TrieWalker<C>,
+        hashed_storage_cursor: H,
+        hashed_address: B256,
+    ) -> Self {
+        Self {
+            walker,
+            hashed_storage_cursor,
+            hashed_address,
+            current_walker_key_checked: false,
+            current_hashed_entry: None,
+        }
+    }
+}
+
+impl<C, H> StorageNodeIter<C, H>
+where
+    C: TrieCursor,
+    H: HashedStorageCursor,
+{
+    /// Return the next storage trie node to be added to the hash builder.
+    ///
+    /// Returns the nodes using this algorithm:
+    /// 1. Return the current intermediate branch node if it hasn't been updated.
+    /// 2. Advance the trie walker to the next intermediate branch node and retrieve next
+    ///    unprocessed key.
+    /// 3. Reposition the hashed storage cursor on the next unprocessed key.
+    /// 4. Return every hashed storage entry up to the key of the current intermediate branch node.
+    /// 5. Repeat.
+    pub(crate) fn try_next(&mut self) -> Result<Option<StorageNode>, StorageRootError> {
+        loop {
+            if let Some(key) = self.walker.key() {
+                if !self.current_walker_key_checked {
+                    self.current_walker_key_checked = true;
+                    if self.walker.can_skip_current_node {
+                        return Ok(Some(StorageNode::Branch(TrieBranchNode::new(
+                            key,
+                            self.walker.hash().unwrap(),
+                            self.walker.children_are_in_trie(),
+                        ))))
+                    }
+                }
+            }
+
+            if let Some(StorageEntry { key: hashed_key, value }) = self.current_hashed_entry.take()
+            {
+                if self.walker.key().map_or(false, |key| key < Nibbles::unpack(hashed_key)) {
+                    self.current_walker_key_checked = false;
+                    continue
+                }
+
+                self.current_hashed_entry = self.hashed_storage_cursor.next()?;
+                return Ok(Some(StorageNode::Leaf(hashed_key, value)))
+            }
+
+            let Some(seek_key) = self.walker.next_unprocessed_key() else { break };
+            self.current_hashed_entry =
+                self.hashed_storage_cursor.seek(self.hashed_address, seek_key)?;
+            self.walker.advance()?;
+        }
+
+        Ok(None)
+    }
+}

crates/trie/src/progress.rs

@@ -1,9 +1,5 @@
 use crate::{trie_cursor::CursorSubNode, updates::TrieUpdates};
-use reth_primitives::{
-    stage::MerkleCheckpoint,
-    trie::{hash_builder::HashBuilder, Nibbles},
-    B256,
-};
+use reth_primitives::{stage::MerkleCheckpoint, trie::hash_builder::HashBuilder, B256};
 
 /// The progress of the state root computation.
 #[derive(Debug)]
@@ -24,8 +20,6 @@ pub struct IntermediateStateRootState {
     pub walker_stack: Vec<CursorSubNode>,
     /// The last hashed account key processed.
     pub last_account_key: B256,
-    /// The last walker key processed.
-    pub last_walker_key: Nibbles,
 }
 
 impl From<MerkleCheckpoint> for IntermediateStateRootState {
@@ -34,7 +28,6 @@ impl From<MerkleCheckpoint> for IntermediateStateRootState {
             hash_builder: HashBuilder::from(value.state),
             walker_stack: value.walker_stack.into_iter().map(CursorSubNode::from).collect(),
             last_account_key: value.last_account_key,
-            last_walker_key: Nibbles::from_hex(value.last_walker_key),
         }
     }
 }

crates/trie/src/proof.rs

@@ -1,10 +1,11 @@
 use crate::{
     account::EthAccount,
-    hashed_cursor::{HashedAccountCursor, HashedCursorFactory, HashedStorageCursor},
+    hashed_cursor::{HashedCursorFactory, HashedStorageCursor},
+    node_iter::{AccountNode, AccountNodeIter, StorageNode, StorageNodeIter},
     prefix_set::PrefixSetMut,
     trie_cursor::{AccountTrieCursor, StorageTrieCursor},
     walker::TrieWalker,
-    StorageRootError,
+    StateRootError, StorageRootError,
 };
 use alloy_rlp::{BufMut, Encodable};
 use reth_db::{tables, transaction::DbTx};
@@ -12,7 +13,7 @@ use reth_primitives::{
     keccak256,
     proofs::EMPTY_ROOT,
     trie::{AccountProof, HashBuilder, Nibbles, StorageProof},
-    Address, StorageEntry, B256,
+    Address, B256,
 };
 
 /// A struct for generating merkle proofs.
@@ -45,65 +46,46 @@ where
         &self,
         address: Address,
         slots: &[B256],
-    ) -> Result<AccountProof, StorageRootError> {
+    ) -> Result<AccountProof, StateRootError> {
         let target_hashed_address = keccak256(address);
         let target_nibbles = Nibbles::unpack(target_hashed_address);
         let mut account_proof = AccountProof::new(address);
 
-        let mut trie_cursor =
-            AccountTrieCursor::new(self.tx.cursor_read::<tables::AccountsTrie>()?);
-        let mut hashed_account_cursor = self.hashed_cursor_factory.hashed_account_cursor()?;
+        let hashed_account_cursor = self.hashed_cursor_factory.hashed_account_cursor()?;
+        let trie_cursor = AccountTrieCursor::new(self.tx.cursor_read::<tables::AccountsTrie>()?);
 
         // Create the walker.
         let mut prefix_set = PrefixSetMut::default();
         prefix_set.insert(target_nibbles.clone());
-        let mut walker = TrieWalker::new(&mut trie_cursor, prefix_set.freeze());
+        let walker = TrieWalker::new(trie_cursor, prefix_set.freeze());
 
         // Create a hash builder to rebuild the root node since it is not available in the database.
         let mut hash_builder =
             HashBuilder::default().with_proof_retainer(Vec::from([target_nibbles.clone()]));
 
         let mut account_rlp = Vec::with_capacity(128);
-        while let Some(key) = walker.key() {
-            if walker.can_skip_current_node {
-                let value = walker.hash().unwrap();
-                let is_in_db_trie = walker.children_are_in_trie();
-                hash_builder.add_branch(key.clone(), value, is_in_db_trie);
-            }
-
-            let seek_key = match walker.next_unprocessed_key() {
-                Some(key) => key,
-                None => break, // no more keys
-            };
-
-            let next_key = walker.advance()?;
-            let mut next_account_entry = hashed_account_cursor.seek(seek_key)?;
-            while let Some((hashed_address, account)) = next_account_entry {
-                let account_nibbles = Nibbles::unpack(hashed_address);
-
-                if let Some(ref key) = next_key {
-                    if key < &account_nibbles {
-                        break
-                    }
+        let mut account_node_iter = AccountNodeIter::new(walker, hashed_account_cursor);
+        while let Some(account_node) = account_node_iter.try_next()? {
+            match account_node {
+                AccountNode::Branch(node) => {
+                    hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
                 }
+                AccountNode::Leaf(hashed_address, account) => {
+                    let storage_root = if hashed_address == target_hashed_address {
+                        let (storage_root, storage_proofs) =
+                            self.storage_root_with_proofs(hashed_address, slots)?;
+                        account_proof.set_account(account, storage_root, storage_proofs);
+                        storage_root
+                    } else {
+                        self.storage_root(hashed_address)?
+                    };
 
-                let storage_root = if hashed_address == target_hashed_address {
-                    let (storage_root, storage_proofs) =
-                        self.storage_root_with_proofs(hashed_address, slots)?;
-                    account_proof.set_account(account, storage_root, storage_proofs);
-                    storage_root
-                } else {
-                    self.storage_root(hashed_address)?
-                };
+                    account_rlp.clear();
+                    let account = EthAccount::from(account).with_storage_root(storage_root);
+                    account.encode(&mut account_rlp as &mut dyn BufMut);
 
-                account_rlp.clear();
-                let account = EthAccount::from(account).with_storage_root(storage_root);
-                account.encode(&mut &mut account_rlp as &mut dyn BufMut);
-
-                hash_builder.add_leaf(account_nibbles, &account_rlp);
-
-                // Move the next account entry
-                next_account_entry = hashed_account_cursor.next()?;
+                    hash_builder.add_leaf(Nibbles::unpack(hashed_address), &account_rlp);
+                }
             }
         }
 
@@ -129,11 +111,6 @@ where
     ) -> Result<(B256, Vec<StorageProof>), StorageRootError> {
         let mut hashed_storage_cursor = self.hashed_cursor_factory.hashed_storage_cursor()?;
 
-        let mut trie_cursor = StorageTrieCursor::new(
-            self.tx.cursor_dup_read::<tables::StoragesTrie>()?,
-            hashed_address,
-        );
-
         let mut proofs = slots.iter().copied().map(StorageProof::new).collect::<Vec<_>>();
 
         // short circuit on empty storage
@@ -143,52 +120,41 @@ where
 
         let target_nibbles = proofs.iter().map(|p| p.nibbles.clone()).collect::<Vec<_>>();
         let prefix_set = PrefixSetMut::from(target_nibbles.clone()).freeze();
-        let mut walker = TrieWalker::new(&mut trie_cursor, prefix_set);
+        let trie_cursor = StorageTrieCursor::new(
+            self.tx.cursor_dup_read::<tables::StoragesTrie>()?,
+            hashed_address,
+        );
+        let walker = TrieWalker::new(trie_cursor, prefix_set);
 
         let mut hash_builder = HashBuilder::default().with_proof_retainer(target_nibbles);
-        while let Some(key) = walker.key() {
-            if walker.can_skip_current_node {
-                hash_builder.add_branch(key, walker.hash().unwrap(), walker.children_are_in_trie());
-            }
-
-            let seek_key = match walker.next_unprocessed_key() {
-                Some(key) => key,
-                None => break, // no more keys
-            };
-
-            let next_key = walker.advance()?;
-            let mut storage = hashed_storage_cursor.seek(hashed_address, seek_key)?;
-            while let Some(StorageEntry { key: hashed_key, value }) = storage {
-                let hashed_key_nibbles = Nibbles::unpack(hashed_key);
-                if let Some(ref key) = next_key {
-                    if key < &hashed_key_nibbles {
-                        break
+        let mut storage_node_iter =
+            StorageNodeIter::new(walker, hashed_storage_cursor, hashed_address);
+        while let Some(node) = storage_node_iter.try_next()? {
+            match node {
+                StorageNode::Branch(node) => {
+                    hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
+                }
+                StorageNode::Leaf(hashed_slot, value) => {
+                    let nibbles = Nibbles::unpack(hashed_slot);
+                    if let Some(proof) = proofs.iter_mut().find(|proof| proof.nibbles == nibbles) {
+                        proof.set_value(value);
                     }
+                    hash_builder.add_leaf(nibbles, alloy_rlp::encode_fixed_size(&value).as_ref());
                 }
-
-                if let Some(proof) =
-                    proofs.iter_mut().find(|proof| proof.nibbles == hashed_key_nibbles)
-                {
-                    proof.set_value(value);
-                }
-
-                hash_builder
-                    .add_leaf(hashed_key_nibbles, alloy_rlp::encode_fixed_size(&value).as_ref());
-                storage = hashed_storage_cursor.next()?;
             }
         }
 
         let root = hash_builder.root();
 
-        let proof_nodes = hash_builder.take_proofs();
+        let all_proof_nodes = hash_builder.take_proofs();
         for proof in proofs.iter_mut() {
-            proof.set_proof(
-                proof_nodes
-                    .iter()
-                    .filter(|(path, _)| proof.nibbles.starts_with(path))
-                    .map(|(_, node)| node.clone())
-                    .collect(),
-            );
+            // Iterate over all proof nodes and find the matching ones.
+            // The filtered results are guaranteed to be in order.
+            let matching_proof_nodes = all_proof_nodes
+                .iter()
+                .filter(|(path, _)| proof.nibbles.starts_with(path))
+                .map(|(_, node)| node.clone());
+            proof.set_proof(matching_proof_nodes.collect());
         }
 
         Ok((root, proofs))

crates/trie/src/trie.rs

@@ -1,6 +1,7 @@
 use crate::{
     account::EthAccount,
-    hashed_cursor::{HashedAccountCursor, HashedCursorFactory, HashedStorageCursor},
+    hashed_cursor::{HashedCursorFactory, HashedStorageCursor},
+    node_iter::{AccountNode, AccountNodeIter, StorageNode, StorageNodeIter},
     prefix_set::{PrefixSet, PrefixSetLoader, PrefixSetMut},
     progress::{IntermediateStateRootState, StateRootProgress},
     trie_cursor::{AccountTrieCursor, StorageTrieCursor},
@@ -8,13 +9,13 @@ use crate::{
     walker::TrieWalker,
     StateRootError, StorageRootError,
 };
-use alloy_rlp::Encodable;
+use alloy_rlp::{BufMut, Encodable};
 use reth_db::{tables, transaction::DbTx};
 use reth_primitives::{
     keccak256,
     proofs::EMPTY_ROOT,
     trie::{HashBuilder, Nibbles},
-    Address, BlockNumber, StorageEntry, B256,
+    Address, BlockNumber, B256,
 };
 use std::{
     collections::{HashMap, HashSet},
@@ -224,136 +225,104 @@ where
         tracing::debug!(target: "loader", "calculating state root");
         let mut trie_updates = TrieUpdates::default();
 
-        let mut hashed_account_cursor = self.hashed_cursor_factory.hashed_account_cursor()?;
-        let mut trie_cursor =
-            AccountTrieCursor::new(self.tx.cursor_read::<tables::AccountsTrie>()?);
+        let hashed_account_cursor = self.hashed_cursor_factory.hashed_account_cursor()?;
+        let trie_cursor = AccountTrieCursor::new(self.tx.cursor_read::<tables::AccountsTrie>()?);
 
-        let (mut walker, mut hash_builder, mut last_account_key, mut last_walker_key) =
-            match self.previous_state {
-                Some(state) => (
-                    TrieWalker::from_stack(
-                        &mut trie_cursor,
-                        state.walker_stack,
-                        self.changed_account_prefixes,
-                    ),
+        let (mut hash_builder, mut account_node_iter) = match self.previous_state {
+            Some(state) => {
+                let walker = TrieWalker::from_stack(
+                    trie_cursor,
+                    state.walker_stack,
+                    self.changed_account_prefixes,
+                );
+                (
                     state.hash_builder,
-                    Some(state.last_account_key),
-                    Some(state.last_walker_key),
-                ),
-                None => (
-                    TrieWalker::new(&mut trie_cursor, self.changed_account_prefixes),
-                    HashBuilder::default(),
-                    None,
-                    None,
-                ),
-            };
+                    AccountNodeIter::new(walker, hashed_account_cursor)
+                        .with_last_account_key(state.last_account_key),
+                )
+            }
+            None => {
+                let walker = TrieWalker::new(trie_cursor, self.changed_account_prefixes);
+                (HashBuilder::default(), AccountNodeIter::new(walker, hashed_account_cursor))
+            }
+        };
 
-        walker.set_updates(retain_updates);
+        account_node_iter.walker.set_updates(retain_updates);
         hash_builder.set_updates(retain_updates);
 
         let mut account_rlp = Vec::with_capacity(128);
         let mut hashed_entries_walked = 0;
-
-        while let Some(key) = last_walker_key.take().or_else(|| walker.key()) {
-            // Take the last account key to make sure we take it into consideration only once.
-            let (next_key, mut next_account_entry) = match last_account_key.take() {
-                // Seek the last processed entry and take the next after.
-                Some(account_key) => {
-                    hashed_account_cursor.seek(account_key)?;
-                    (walker.key(), hashed_account_cursor.next()?)
+        while let Some(node) = account_node_iter.try_next()? {
+            match node {
+                AccountNode::Branch(node) => {
+                    hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
                 }
-                None => {
-                    if walker.can_skip_current_node {
-                        let value = walker.hash().unwrap();
-                        let is_in_db_trie = walker.children_are_in_trie();
-                        hash_builder.add_branch(key.clone(), value, is_in_db_trie);
-                    }
+                AccountNode::Leaf(hashed_address, account) => {
+                    hashed_entries_walked += 1;
 
-                    let seek_key = match walker.next_unprocessed_key() {
-                        Some(key) => key,
-                        None => break, // no more keys
+                    // We assume we can always calculate a storage root without
+                    // OOMing. This opens us up to a potential DOS vector if
+                    // a contract had too many storage entries and they were
+                    // all buffered w/o us returning and committing our intermediate
+                    // progress.
+                    // TODO: We can consider introducing the TrieProgress::Progress/Complete
+                    // abstraction inside StorageRoot, but let's give it a try as-is for now.
+                    let storage_root_calculator = StorageRoot::new_hashed(self.tx, hashed_address)
+                        .with_hashed_cursor_factory(self.hashed_cursor_factory.clone())
+                        .with_changed_prefixes(
+                            self.changed_storage_prefixes
+                                .get(&hashed_address)
+                                .cloned()
+                                .unwrap_or_default(),
+                        );
+
+                    let storage_root = if retain_updates {
+                        let (root, storage_slots_walked, updates) =
+                            storage_root_calculator.root_with_updates()?;
+                        hashed_entries_walked += storage_slots_walked;
+                        trie_updates.extend(updates.into_iter());
+                        root
+                    } else {
+                        storage_root_calculator.root()?
                     };
 
-                    (walker.advance()?, hashed_account_cursor.seek(seek_key)?)
-                }
-            };
+                    let account = EthAccount::from(account).with_storage_root(storage_root);
 
-            while let Some((hashed_address, account)) = next_account_entry {
-                hashed_entries_walked += 1;
-                let account_nibbles = Nibbles::unpack(hashed_address);
+                    account_rlp.clear();
+                    account.encode(&mut account_rlp as &mut dyn BufMut);
 
-                if let Some(ref key) = next_key {
-                    if key < &account_nibbles {
-                        tracing::trace!(target: "loader", "breaking, already detected");
-                        break
+                    hash_builder.add_leaf(Nibbles::unpack(hashed_address), &account_rlp);
+
+                    // Decide if we need to return intermediate progress.
+                    let total_updates_len = trie_updates.len() +
+                        account_node_iter.walker.updates_len() +
+                        hash_builder.updates_len();
+                    if retain_updates && total_updates_len as u64 >= self.threshold {
+                        let (walker_stack, walker_updates) = account_node_iter.walker.split();
+                        let (hash_builder, hash_builder_updates) = hash_builder.split();
+
+                        let state = IntermediateStateRootState {
+                            hash_builder,
+                            walker_stack,
+                            last_account_key: hashed_address,
+                        };
+
+                        trie_updates.extend(walker_updates.into_iter());
+                        trie_updates.extend_with_account_updates(hash_builder_updates);
+
+                        return Ok(StateRootProgress::Progress(
+                            Box::new(state),
+                            hashed_entries_walked,
+                            trie_updates,
+                        ))
                     }
                 }
-
-                // We assume we can always calculate a storage root without
-                // OOMing. This opens us up to a potential DOS vector if
-                // a contract had too many storage entries and they were
-                // all buffered w/o us returning and committing our intermediate
-                // progress.
-                // TODO: We can consider introducing the TrieProgress::Progress/Complete
-                // abstraction inside StorageRoot, but let's give it a try as-is for now.
-                let storage_root_calculator = StorageRoot::new_hashed(self.tx, hashed_address)
-                    .with_hashed_cursor_factory(self.hashed_cursor_factory.clone())
-                    .with_changed_prefixes(
-                        self.changed_storage_prefixes
-                            .get(&hashed_address)
-                            .cloned()
-                            .unwrap_or_default(),
-                    );
-
-                let storage_root = if retain_updates {
-                    let (root, storage_slots_walked, updates) =
-                        storage_root_calculator.root_with_updates()?;
-                    hashed_entries_walked += storage_slots_walked;
-                    trie_updates.extend(updates.into_iter());
-                    root
-                } else {
-                    storage_root_calculator.root()?
-                };
-
-                let account = EthAccount::from(account).with_storage_root(storage_root);
-
-                account_rlp.clear();
-                account.encode(&mut &mut account_rlp);
-
-                hash_builder.add_leaf(account_nibbles, &account_rlp);
-
-                // Decide if we need to return intermediate progress.
-                let total_updates_len =
-                    trie_updates.len() + walker.updates_len() + hash_builder.updates_len();
-                if retain_updates && total_updates_len as u64 >= self.threshold {
-                    let (walker_stack, walker_updates) = walker.split();
-                    let (hash_builder, hash_builder_updates) = hash_builder.split();
-
-                    let state = IntermediateStateRootState {
-                        hash_builder,
-                        walker_stack,
-                        last_walker_key: key,
-                        last_account_key: hashed_address,
-                    };
-
-                    trie_updates.extend(walker_updates.into_iter());
-                    trie_updates.extend_with_account_updates(hash_builder_updates);
-
-                    return Ok(StateRootProgress::Progress(
-                        Box::new(state),
-                        hashed_entries_walked,
-                        trie_updates,
-                    ))
-                }
-
-                // Move the next account entry
-                next_account_entry = hashed_account_cursor.next()?;
             }
         }
 
         let root = hash_builder.root();
 
-        let (_, walker_updates) = walker.split();
+        let (_, walker_updates) = account_node_iter.walker.split();
         let (_, hash_builder_updates) = hash_builder.split();
 
         trie_updates.extend(walker_updates.into_iter());
@@ -464,14 +433,8 @@ where
         retain_updates: bool,
     ) -> Result<(B256, usize, TrieUpdates), StorageRootError> {
         tracing::debug!(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 trie_cursor = StorageTrieCursor::new(
-            self.tx.cursor_dup_read::<tables::StoragesTrie>()?,
-            self.hashed_address,
-        );
-
         // short circuit on empty storage
         if hashed_storage_cursor.is_storage_empty(self.hashed_address)? {
             return Ok((
@@ -481,43 +444,37 @@ where
             ))
         }
 
-        let mut walker = TrieWalker::new(&mut trie_cursor, self.changed_prefixes.clone())
+        let trie_cursor = StorageTrieCursor::new(
+            self.tx.cursor_dup_read::<tables::StoragesTrie>()?,
+            self.hashed_address,
+        );
+        let walker = TrieWalker::new(trie_cursor, self.changed_prefixes.clone())
             .with_updates(retain_updates);
 
         let mut hash_builder = HashBuilder::default().with_updates(retain_updates);
 
         let mut storage_slots_walked = 0;
-        while let Some(key) = walker.key() {
-            if walker.can_skip_current_node {
-                hash_builder.add_branch(key, walker.hash().unwrap(), walker.children_are_in_trie());
-            }
-
-            let seek_key = match walker.next_unprocessed_key() {
-                Some(key) => key,
-                None => break, // no more keys
-            };
-
-            let next_key = walker.advance()?;
-            let mut storage = hashed_storage_cursor.seek(self.hashed_address, seek_key)?;
-            while let Some(StorageEntry { key: hashed_key, value }) = storage {
-                storage_slots_walked += 1;
-
-                let storage_key_nibbles = Nibbles::unpack(hashed_key);
-                if let Some(ref key) = next_key {
-                    if key < &storage_key_nibbles {
-                        break
-                    }
+        let mut storage_node_iter =
+            StorageNodeIter::new(walker, hashed_storage_cursor, self.hashed_address);
+        while let Some(node) = storage_node_iter.try_next()? {
+            match node {
+                StorageNode::Branch(node) => {
+                    hash_builder.add_branch(node.key, node.value, node.children_are_in_trie);
+                }
+                StorageNode::Leaf(hashed_slot, value) => {
+                    storage_slots_walked += 1;
+                    hash_builder.add_leaf(
+                        Nibbles::unpack(hashed_slot),
+                        alloy_rlp::encode_fixed_size(&value).as_ref(),
+                    );
                 }
-                hash_builder
-                    .add_leaf(storage_key_nibbles, alloy_rlp::encode_fixed_size(&value).as_ref());
-                storage = hashed_storage_cursor.next()?;
             }
         }
 
         let root = hash_builder.root();
 
         let (_, hash_builder_updates) = hash_builder.split();
-        let (_, walker_updates) = walker.split();
+        let (_, walker_updates) = storage_node_iter.walker.split();
 
         let mut trie_updates = TrieUpdates::default();
         trie_updates.extend(walker_updates.into_iter());
@@ -529,7 +486,6 @@ where
 }
 
 #[cfg(test)]
-#[allow(clippy::mutable_key_type)]
 mod tests {
     use super::*;
     use crate::test_utils::{
@@ -548,7 +504,7 @@ mod tests {
         keccak256,
         proofs::triehash::KeccakHasher,
         trie::{BranchNodeCompact, TrieMask},
-        Account, Address, B256, MAINNET, U256,
+        Account, Address, StorageEntry, B256, MAINNET, U256,
     };
     use reth_provider::{DatabaseProviderRW, ProviderFactory};
     use std::{collections::BTreeMap, ops::Mul, str::FromStr};
@@ -621,7 +577,6 @@ mod tests {
     }
 
     #[test]
-    // TODO: Try to find the edge case by creating some more very complex trie.
     fn branch_node_child_changes() {
         incremental_vs_full_root(
             &[

crates/trie/src/trie_cursor/account_cursor.rs

@@ -14,20 +14,22 @@ impl<C> AccountTrieCursor<C> {
     }
 }
 
-impl<C> TrieCursor<StoredNibbles> for AccountTrieCursor<C>
+impl<C> TrieCursor for AccountTrieCursor<C>
 where
     C: DbCursorRO<tables::AccountsTrie>,
 {
+    type Key = StoredNibbles;
+
     fn seek_exact(
         &mut self,
-        key: StoredNibbles,
+        key: Self::Key,
     ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError> {
         Ok(self.0.seek_exact(key)?.map(|value| (value.0.inner.to_vec(), value.1)))
     }
 
     fn seek(
         &mut self,
-        key: StoredNibbles,
+        key: Self::Key,
     ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError> {
         Ok(self.0.seek(key)?.map(|value| (value.0.inner.to_vec(), value.1)))
     }

crates/trie/src/trie_cursor/mod.rs

@@ -1,5 +1,5 @@
 use crate::updates::TrieKey;
-use reth_db::{table::Key, DatabaseError};
+use reth_db::DatabaseError;
 use reth_primitives::trie::BranchNodeCompact;
 
 mod account_cursor;
@@ -11,13 +11,22 @@ pub use self::{
 };
 
 /// A cursor for navigating a trie that works with both Tables and DupSort tables.
-pub trait TrieCursor<K: Key> {
+#[auto_impl::auto_impl(&mut)]
+pub trait TrieCursor {
+    /// The key type of the cursor.
+    type Key: From<Vec<u8>>;
+
     /// Move the cursor to the key and return if it is an exact match.
-    fn seek_exact(&mut self, key: K)
-        -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError>;
+    fn seek_exact(
+        &mut self,
+        key: Self::Key,
+    ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError>;
 
     /// Move the cursor to the key and return a value matching of greater than the key.
-    fn seek(&mut self, key: K) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError>;
+    fn seek(
+        &mut self,
+        key: Self::Key,
+    ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError>;
 
     /// Get the current entry.
     fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError>;

crates/trie/src/trie_cursor/storage_cursor.rs

@@ -24,13 +24,15 @@ impl<C> StorageTrieCursor<C> {
     }
 }
 
-impl<C> TrieCursor<StoredNibblesSubKey> for StorageTrieCursor<C>
+impl<C> TrieCursor for StorageTrieCursor<C>
 where
     C: DbDupCursorRO<tables::StoragesTrie> + DbCursorRO<tables::StoragesTrie>,
 {
+    type Key = StoredNibblesSubKey;
+
     fn seek_exact(
         &mut self,
-        key: StoredNibblesSubKey,
+        key: Self::Key,
     ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError> {
         Ok(self
             .cursor
@@ -41,7 +43,7 @@ where
 
     fn seek(
         &mut self,
-        key: StoredNibblesSubKey,
+        key: Self::Key,
     ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError> {
         Ok(self
             .cursor

crates/trie/src/updates.rs

@@ -77,7 +77,6 @@ impl TrieUpdates {
     }
 
     /// Extend the updates with account trie updates.
-    #[allow(clippy::mutable_key_type)]
     pub fn extend_with_account_updates(&mut self, updates: HashMap<Nibbles, BranchNodeCompact>) {
         self.extend(updates.into_iter().map(|(nibbles, node)| {
             (TrieKey::AccountNode(nibbles.hex_data.to_vec().into()), TrieOp::Update(node))
@@ -85,7 +84,6 @@ impl TrieUpdates {
     }
 
     /// Extend the updates with storage trie updates.
-    #[allow(clippy::mutable_key_type)]
     pub fn extend_with_storage_updates(
         &mut self,
         hashed_address: B256,

crates/trie/src/walker.rs

@@ -3,19 +3,19 @@ use crate::{
     trie_cursor::{CursorSubNode, TrieCursor},
     updates::TrieUpdates,
 };
-use reth_db::{table::Key, DatabaseError};
+use reth_db::DatabaseError;
 use reth_primitives::{
     trie::{BranchNodeCompact, Nibbles},
     B256,
 };
-use std::marker::PhantomData;
 
 /// `TrieWalker` is a structure that enables traversal of a Merkle trie.
-/// It allows moving through the trie in a depth-first manner, skipping certain branches if the .
+/// It allows moving through the trie in a depth-first manner, skipping certain branches
+/// if they have not changed.
 #[derive(Debug)]
-pub struct TrieWalker<'a, K, C> {
+pub struct TrieWalker<C> {
     /// A mutable reference to a trie cursor instance used for navigating the trie.
-    pub cursor: &'a mut C,
+    pub cursor: C,
     /// A vector containing the trie nodes that have been visited.
     pub stack: Vec<CursorSubNode>,
     /// A flag indicating whether the current node can be skipped when traversing the trie. This
@@ -26,12 +26,11 @@ pub struct TrieWalker<'a, K, C> {
     pub changes: PrefixSet,
     /// The trie updates to be applied to the trie.
     trie_updates: Option<TrieUpdates>,
-    __phantom: PhantomData<K>,
 }
 
-impl<'a, K: Key + From<Vec<u8>>, C: TrieCursor<K>> TrieWalker<'a, K, C> {
+impl<C: TrieCursor> TrieWalker<C> {
     /// Constructs a new TrieWalker, setting up the initial state of the stack and cursor.
-    pub fn new(cursor: &'a mut C, changes: PrefixSet) -> Self {
+    pub fn new(cursor: C, changes: PrefixSet) -> Self {
         // Initialize the walker with a single empty stack element.
         let mut this = Self {
             cursor,
@@ -39,7 +38,6 @@ impl<'a, K: Key + From<Vec<u8>>, C: TrieCursor<K>> TrieWalker<'a, K, C> {
             stack: vec![CursorSubNode::default()],
             can_skip_current_node: false,
             trie_updates: None,
-            __phantom: PhantomData,
         };
 
         // Set up the root node of the trie in the stack, if it exists.
@@ -53,15 +51,9 @@ impl<'a, K: Key + From<Vec<u8>>, C: TrieCursor<K>> TrieWalker<'a, K, C> {
     }
 
     /// Constructs a new TrieWalker from existing stack and a cursor.
-    pub fn from_stack(cursor: &'a mut C, stack: Vec<CursorSubNode>, changes: PrefixSet) -> Self {
-        let mut this = Self {
-            cursor,
-            changes,
-            stack,
-            can_skip_current_node: false,
-            trie_updates: None,
-            __phantom: PhantomData,
-        };
+    pub fn from_stack(cursor: C, stack: Vec<CursorSubNode>, changes: PrefixSet) -> Self {
+        let mut this =
+            Self { cursor, changes, stack, can_skip_current_node: false, trie_updates: None };
         this.update_skip_node();
         this
     }
@@ -255,7 +247,6 @@ impl<'a, K: Key + From<Vec<u8>>, C: TrieCursor<K>> TrieWalker<'a, K, C> {
 
 #[cfg(test)]
 mod tests {
-
     use super::*;
     use crate::{
         prefix_set::PrefixSetMut,
@@ -316,10 +307,9 @@ mod tests {
         test_cursor(storage_trie, &expected);
     }
 
-    fn test_cursor<K, T>(mut trie: T, expected: &[Vec<u8>])
+    fn test_cursor<T>(mut trie: T, expected: &[Vec<u8>])
     where
-        K: Key + From<Vec<u8>>,
-        T: TrieCursor<K>,
+        T: TrieCursor,
     {
         let mut walker = TrieWalker::new(&mut trie, Default::default());
         assert!(walker.key().unwrap().is_empty());