chore(trie): store only deleted keys in TrieWalker

crates/trie/parallel/src/async_root.rs

@@ -132,7 +132,7 @@ where
             trie_cursor_factory.account_trie_cursor().map_err(ProviderError::Database)?,
             prefix_sets.account_prefix_set,
         )
-        .with_updates(retain_updates);
+        .with_deletions_retained(retain_updates);
         let mut account_node_iter = TrieNodeIter::new(
             walker,
             hashed_cursor_factory.hashed_account_cursor().map_err(ProviderError::Database)?,

crates/trie/parallel/src/parallel_root.rs

@@ -116,7 +116,7 @@ where
             trie_cursor_factory.account_trie_cursor().map_err(ProviderError::Database)?,
             prefix_sets.account_prefix_set,
         )
-        .with_updates(retain_updates);
+        .with_deletions_retained(retain_updates);
         let mut account_node_iter = TrieNodeIter::new(
             walker,
             hashed_cursor_factory.hashed_account_cursor().map_err(ProviderError::Database)?,

crates/trie/trie/src/trie.rs

@@ -221,15 +221,15 @@ where
                     state.walker_stack,
                     self.prefix_sets.account_prefix_set,
                 )
-                .with_updates(retain_updates);
+                .with_deletions_retained(retain_updates);
                 let node_iter = TrieNodeIter::new(walker, hashed_account_cursor)
                     .with_last_hashed_key(state.last_account_key);
                 (hash_builder, node_iter)
             }
             None => {
                 let hash_builder = HashBuilder::default().with_updates(retain_updates);
                 let walker = TrieWalker::new(trie_cursor, self.prefix_sets.account_prefix_set)
-                    .with_updates(retain_updates);
+                    .with_deletions_retained(retain_updates);
                 let node_iter = TrieNodeIter::new(walker, hashed_account_cursor);
                 (hash_builder, node_iter)
             }
@@ -286,10 +286,10 @@ where
 
                     // Decide if we need to return intermediate progress.
                     let total_updates_len = trie_updates.len() +
-                        account_node_iter.walker.updates_len() +
+                        account_node_iter.walker.deleted_keys_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 (walker_stack, walker_deleted_keys) = account_node_iter.walker.split();
                         let (hash_builder, hash_builder_updates) = hash_builder.split();
 
                         let state = IntermediateStateRootState {
@@ -298,7 +298,9 @@ where
                             last_account_key: hashed_address,
                         };
 
-                        trie_updates.extend(walker_updates);
+                        trie_updates.extend(
+                            walker_deleted_keys.into_iter().map(|key| (key, TrieOp::Delete)),
+                        );
                         trie_updates.extend_with_account_updates(hash_builder_updates);
 
                         return Ok(StateRootProgress::Progress(
@@ -492,7 +494,8 @@ where
 
         let mut tracker = TrieTracker::default();
         let trie_cursor = self.trie_cursor_factory.storage_trie_cursor(self.hashed_address)?;
-        let walker = TrieWalker::new(trie_cursor, self.prefix_set).with_updates(retain_updates);
+        let walker =
+            TrieWalker::new(trie_cursor, self.prefix_set).with_deletions_retained(retain_updates);
 
         let mut hash_builder = HashBuilder::default().with_updates(retain_updates);
 

crates/trie/trie/src/updates.rs

@@ -110,18 +110,6 @@ impl IntoIterator for TrieUpdates {
 }
 
 impl TrieUpdates {
-    /// Schedule a delete operation on a trie key.
-    ///
-    /// # Panics
-    ///
-    /// If the key already exists and the operation is an update.
-    pub fn schedule_delete(&mut self, key: TrieKey) {
-        let existing = self.trie_operations.insert(key, TrieOp::Delete);
-        if let Some(op) = existing {
-            assert!(!op.is_update(), "Tried to delete a node that was already updated");
-        }
-    }
-
     /// Extend the updates with trie updates.
     pub fn extend(&mut self, updates: impl IntoIterator<Item = (TrieKey, TrieOp)>) {
         self.trie_operations.extend(updates);
@@ -144,8 +132,8 @@ impl TrieUpdates {
         destroyed_accounts: HashSet<B256>,
     ) {
         // Add updates from trie walker.
-        let (_, walker_updates) = walker.split();
-        self.extend(walker_updates);
+        let (_, deleted_keys) = walker.split();
+        self.extend(deleted_keys.into_iter().map(|key| (key, TrieOp::Delete)));
 
         // Add account node updates from hash builder.
         let (_, hash_builder_updates) = hash_builder.split();
@@ -165,8 +153,8 @@ impl TrieUpdates {
         hash_builder: HashBuilder,
     ) {
         // Add updates from trie walker.
-        let (_, walker_updates) = walker.split();
-        self.extend(walker_updates);
+        let (_, deleted_keys) = walker.split();
+        self.extend(deleted_keys.into_iter().map(|key| (key, TrieOp::Delete)));
 
         // Add storage node updates from hash builder.
         let (_, hash_builder_updates) = hash_builder.split();

crates/trie/trie/src/walker.rs

@@ -1,11 +1,12 @@
 use crate::{
     prefix_set::PrefixSet,
     trie_cursor::{CursorSubNode, TrieCursor},
-    updates::TrieUpdates,
+    updates::TrieKey,
     BranchNodeCompact, Nibbles,
 };
 use reth_db::DatabaseError;
 use reth_primitives::B256;
+use std::collections::HashSet;
 
 /// `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
@@ -22,36 +23,31 @@ pub struct TrieWalker<C> {
     pub can_skip_current_node: bool,
     /// A `PrefixSet` representing the changes to be applied to the trie.
     pub changes: PrefixSet,
-    /// The trie updates to be applied to the trie.
-    trie_updates: Option<TrieUpdates>,
+    /// The retained trie node keys that need to be deleted.
+    deleted_keys: Option<HashSet<TrieKey>>,
 }
 
 impl<C> TrieWalker<C> {
     /// Constructs a new `TrieWalker` from existing stack and a cursor.
     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 };
+            Self { cursor, changes, stack, can_skip_current_node: false, deleted_keys: None };
         this.update_skip_node();
         this
     }
 
     /// Sets the flag whether the trie updates should be stored.
-    pub fn with_updates(mut self, retain_updates: bool) -> Self {
-        self.set_updates(retain_updates);
-        self
-    }
-
-    /// Sets the flag whether the trie updates should be stored.
-    pub fn set_updates(&mut self, retain_updates: bool) {
-        if retain_updates {
-            self.trie_updates = Some(TrieUpdates::default());
+    pub fn with_deletions_retained(mut self, retained: bool) -> Self {
+        if retained {
+            self.deleted_keys = Some(HashSet::default());
         }
+        self
     }
 
     /// Split the walker into stack and trie updates.
-    pub fn split(mut self) -> (Vec<CursorSubNode>, TrieUpdates) {
-        let trie_updates = self.trie_updates.take();
-        (self.stack, trie_updates.unwrap_or_default())
+    pub fn split(mut self) -> (Vec<CursorSubNode>, HashSet<TrieKey>) {
+        let keys = self.deleted_keys.take();
+        (self.stack, keys.unwrap_or_default())
     }
 
     /// Prints the current stack of trie nodes.
@@ -63,9 +59,9 @@ impl<C> TrieWalker<C> {
         println!("====================== END STACK ======================\n");
     }
 
-    /// The current length of the trie updates.
-    pub fn updates_len(&self) -> usize {
-        self.trie_updates.as_ref().map(|u| u.len()).unwrap_or(0)
+    /// The current length of the deleted keys.
+    pub fn deleted_keys_len(&self) -> usize {
+        self.deleted_keys.as_ref().map_or(0, |u| u.len())
     }
 
     /// Returns the current key in the trie.
@@ -117,7 +113,7 @@ impl<C: TrieCursor> TrieWalker<C> {
             changes,
             stack: vec![CursorSubNode::default()],
             can_skip_current_node: false,
-            trie_updates: None,
+            deleted_keys: None,
         };
 
         // Set up the root node of the trie in the stack, if it exists.
@@ -193,8 +189,8 @@ impl<C: TrieCursor> TrieWalker<C> {
         // Delete the current node if it's included in the prefix set or it doesn't contain the root
         // hash.
         if !self.can_skip_current_node || nibble != -1 {
-            if let Some((updates, key)) = self.trie_updates.as_mut().zip(self.cursor.current()?) {
-                updates.schedule_delete(key);
+            if let Some((keys, key)) = self.deleted_keys.as_mut().zip(self.cursor.current()?) {
+                keys.insert(key);
             }
         }