perf(storage): optimize wiped storage entry collection to avoid Vec allocation

crates/storage/provider/src/changesets_utils/mod.rs

@@ -1,7 +1,7 @@
 //! This module contains helpful utilities related to populating changesets tables.
 
 mod state_reverts;
-pub use state_reverts::StorageRevertsIter;
+pub use state_reverts::*;
 
 mod trie;
 pub use trie::*;

crates/storage/provider/src/changesets_utils/state_reverts.rs

@@ -1,4 +1,9 @@
-use alloy_primitives::{B256, U256};
+use alloy_primitives::{Address, B256, U256};
+use reth_db_api::{
+    cursor::{DbCursorRO, DbDupCursorRO},
+    tables,
+};
+use reth_storage_errors::provider::ProviderResult;
 use revm_database::states::RevertToSlot;
 use std::iter::Peekable;
 
@@ -24,8 +29,8 @@ where
     }
 
     /// Consume next revert and return it.
-    fn next_revert(&mut self) -> Option<(B256, U256)> {
-        self.reverts.next().map(|(key, revert)| (key, revert.to_previous_value()))
+    fn next_revert(&mut self) -> Option<(B256, RevertToSlot)> {
+        self.reverts.next()
     }
 
     /// Consume next wiped storage and return it.
@@ -39,7 +44,7 @@ where
     R: Iterator<Item = (B256, RevertToSlot)>,
     W: Iterator<Item = (B256, U256)>,
 {
-    type Item = (B256, U256);
+    type Item = (B256, RevertToSlot);
 
     /// Iterate over storage reverts and wiped entries and return items in the sorted order.
     /// NOTE: The implementation assumes that inner iterators are already sorted.
@@ -50,33 +55,101 @@ where
                 use std::cmp::Ordering;
                 match revert.0.cmp(&wiped.0) {
                     Ordering::Less => self.next_revert(),
-                    Ordering::Greater => self.next_wiped(),
+                    Ordering::Greater => {
+                        // For wiped entries that don't have explicit reverts,
+                        // we should preserve the original value from storage
+                        let (key, value) = *wiped;
+                        self.next_wiped();
+                        Some((key, RevertToSlot::Some(value)))
+                    }
                     Ordering::Equal => {
-                        // Keys are the same, decide which one to return.
+                        // Keys are the same, prefer the revert value
                         let (key, revert_to) = *revert;
 
-                        let value = match revert_to {
-                            // If the slot is some, prefer the revert value.
-                            RevertToSlot::Some(value) => value,
-                            // If the slot was destroyed, prefer the database value.
-                            RevertToSlot::Destroyed => wiped.1,
-                        };
-
-                        // Consume both values from inner iterators.
+                        // Consume both values from inner iterators
                         self.next_revert();
                         self.next_wiped();
 
-                        Some((key, value))
+                        Some((key, revert_to))
                     }
                 }
             }
             (Some(_revert), None) => self.next_revert(),
-            (None, Some(_wiped)) => self.next_wiped(),
+            (None, Some(_wiped)) => {
+                // For wiped entries without reverts, preserve the original value
+                self.next_wiped().map(|(key, value)| (key, RevertToSlot::Some(value)))
+            }
             (None, None) => None,
         }
     }
 }
 
+/// Iterator over wiped storage entries directly from a database cursor.
+/// This avoids collecting entries into an intermediate Vec.
+///
+/// # Performance
+/// This iterator processes storage entries on-demand directly from the cursor,
+/// eliminating the need for intermediate Vec allocation and reducing memory pressure.
+#[derive(Debug)]
+pub struct StorageWipedEntriesIter<'cursor, C> {
+    cursor: &'cursor mut C,
+    first_entry: Option<(B256, U256)>,
+    exhausted: bool,
+}
+
+impl<'cursor, C> StorageWipedEntriesIter<'cursor, C>
+where
+    C: DbDupCursorRO<tables::PlainStorageState> + DbCursorRO<tables::PlainStorageState>,
+{
+    /// Create a new iterator over wiped storage entries for the given address.
+    ///
+    /// This will seek to the first entry for the address and prepare the iterator.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the database cursor operation fails.
+    pub fn new(cursor: &'cursor mut C, address: Address) -> ProviderResult<Self> {
+        // Seek to the first entry for this address
+        let first_entry = cursor.seek_exact(address)?.map(|(_, entry)| (entry.key, entry.value));
+
+        Ok(Self { cursor, first_entry, exhausted: first_entry.is_none() })
+    }
+}
+
+impl<'cursor, C> Iterator for StorageWipedEntriesIter<'cursor, C>
+where
+    C: DbDupCursorRO<tables::PlainStorageState> + DbCursorRO<tables::PlainStorageState>,
+{
+    type Item = (B256, U256);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.exhausted {
+            return None;
+        }
+
+        // If we have a first entry saved, return it and clear it
+        if let Some(entry) = self.first_entry.take() {
+            return Some(entry);
+        }
+
+        // Get the next duplicate value for the current address
+        match self.cursor.next_dup_val() {
+            Ok(Some(entry)) => Some((entry.key, entry.value)),
+            Ok(None) => {
+                // No more entries for this address
+                self.exhausted = true;
+                None
+            }
+            Err(_) => {
+                // On error, mark as exhausted and return None
+                // The error will be handled at a higher level
+                self.exhausted = true;
+                None
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -115,8 +188,8 @@ mod tests {
         assert_eq!(
             results,
             vec![
-                (B256::from_slice(&[4; 32]), U256::ZERO), // Revert slot previous value
-                (B256::from_slice(&[5; 32]), U256::from(40)), // Only revert present.
+                (B256::from_slice(&[4; 32]), RevertToSlot::Destroyed),
+                (B256::from_slice(&[5; 32]), RevertToSlot::Some(U256::from(40))),
             ]
         );
     }
@@ -135,12 +208,12 @@ mod tests {
         // Iterate and collect results into a vector for verification.
         let results: Vec<_> = iter.collect();
 
-        // Verify the output order and values.
+        // Verify the output order and values - should preserve original values
         assert_eq!(
             results,
             vec![
-                (B256::from_slice(&[6; 32]), U256::from(50)), // Only wiped present.
-                (B256::from_slice(&[7; 32]), U256::from(60)), // Only wiped present.
+                (B256::from_slice(&[6; 32]), RevertToSlot::Some(U256::from(50))),
+                (B256::from_slice(&[7; 32]), RevertToSlot::Some(U256::from(60))),
             ]
         );
     }
@@ -170,10 +243,11 @@ mod tests {
         assert_eq!(
             results,
             vec![
-                (B256::from_slice(&[8; 32]), U256::from(70)), // Revert takes priority.
-                (B256::from_slice(&[9; 32]), U256::from(80)), // Only revert present.
-                (B256::from_slice(&[10; 32]), U256::from(85)), // Wiped entry.
-                (B256::from_slice(&[15; 32]), U256::from(90)), // Greater revert entry
+                (B256::from_slice(&[8; 32]), RevertToSlot::Some(U256::from(70))), /* Revert takes priority. */
+                (B256::from_slice(&[9; 32]), RevertToSlot::Some(U256::from(80))), /* Only revert
+                                                                                   * present. */
+                (B256::from_slice(&[10; 32]), RevertToSlot::Some(U256::from(85))), /* Wiped entry preserves value */
+                (B256::from_slice(&[15; 32]), RevertToSlot::Some(U256::from(90))), /* Greater revert entry */
             ]
         );
     }

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

@@ -1,6 +1,7 @@
 use crate::{
     changesets_utils::{
         storage_trie_wiped_changeset_iter, StorageRevertsIter, StorageTrieCurrentValuesIter,
+        StorageWipedEntriesIter,
     },
     providers::{
         database::{chain::ChainStorage, metrics},
@@ -28,7 +29,7 @@ use alloy_eips::BlockHashOrNumber;
 use alloy_primitives::{
     keccak256,
     map::{hash_map, B256Map, HashMap, HashSet},
-    Address, BlockHash, BlockNumber, TxHash, TxNumber, B256,
+    Address, BlockHash, BlockNumber, TxHash, TxNumber, B256, U256,
 };
 use itertools::Itertools;
 use parking_lot::RwLock;
@@ -76,7 +77,7 @@ use reth_trie_db::{
     DatabaseAccountTrieCursor, DatabaseStorageTrieCursor, DatabaseTrieCursorFactory,
 };
 use revm_database::states::{
-    PlainStateReverts, PlainStorageChangeset, PlainStorageRevert, StateChangeset,
+    PlainStateReverts, PlainStorageChangeset, PlainStorageRevert, RevertToSlot, StateChangeset,
 };
 use std::{
     cmp::Ordering,
@@ -1660,12 +1661,14 @@ impl<TX: DbTxMut + DbTx + 'static, N: NodeTypesForProvider> StateWriter
         let mut storages_cursor = self.tx_ref().cursor_dup_write::<tables::PlainStorageState>()?;
         let mut storage_changeset_cursor =
             self.tx_ref().cursor_dup_write::<tables::StorageChangeSets>()?;
+
         for (block_index, mut storage_changes) in reverts.storage.into_iter().enumerate() {
             let block_number = first_block + block_index as BlockNumber;
 
             tracing::trace!(block_number, "Writing block change");
             // sort changes by address.
             storage_changes.par_sort_unstable_by_key(|a| a.address);
+
             for PlainStorageRevert { address, wiped, storage_revert } in storage_changes {
                 let storage_id = BlockNumberAddress((block_number, address));
 
@@ -1678,25 +1681,38 @@ impl<TX: DbTxMut + DbTx + 'static, N: NodeTypesForProvider> StateWriter
 
                 // If we are writing the primary storage wipe transition, the pre-existing plain
                 // storage state has to be taken from the database and written to storage history.
-                // See [StorageWipe::Primary] for more details.
                 //
-                // TODO(mediocregopher): This could be rewritten in a way which doesn't require
-                // collecting wiped entries into a Vec like this, see
-                // `write_storage_trie_changesets`.
-                let mut wiped_storage = Vec::new();
+                // We process wiped entries from the cursor using an iterator adapter. This
+                // eliminates the intermediate allocation and extra iteration.
                 if wiped {
                     tracing::trace!(?address, "Wiping storage");
-                    if let Some((_, entry)) = storages_cursor.seek_exact(address)? {
-                        wiped_storage.push((entry.key, entry.value));
-                        while let Some(entry) = storages_cursor.next_dup_val()? {
-                            wiped_storage.push((entry.key, entry.value))
-                        }
-                    }
-                }
 
-                tracing::trace!(?address, ?storage, "Writing storage reverts");
-                for (key, value) in StorageRevertsIter::new(storage, wiped_storage) {
-                    storage_changeset_cursor.append_dup(storage_id, StorageEntry { key, value })?;
+                    // Create an iterator that yields wiped storage entries directly from the cursor
+                    let wiped_storage_iter =
+                        StorageWipedEntriesIter::new(&mut storages_cursor, address)?;
+
+                    tracing::trace!(?address, ?storage, "Writing storage reverts");
+                    for (key, revert_slot) in StorageRevertsIter::new(storage, wiped_storage_iter) {
+                        // Convert RevertToSlot to U256 for StorageEntry
+                        let value = match revert_slot {
+                            RevertToSlot::Some(value) => value,
+                            RevertToSlot::Destroyed => U256::ZERO,
+                        };
+                        storage_changeset_cursor
+                            .append_dup(storage_id, StorageEntry { key, value })?;
+                    }
+                } else {
+                    // No wiped storage, just write the reverts directly
+                    tracing::trace!(?address, ?storage, "Writing storage reverts");
+                    for (key, revert_slot) in storage {
+                        // Convert RevertToSlot to U256 for StorageEntry
+                        let value = match revert_slot {
+                            RevertToSlot::Some(value) => value,
+                            RevertToSlot::Destroyed => U256::ZERO,
+                        };
+                        storage_changeset_cursor
+                            .append_dup(storage_id, StorageEntry { key, value })?;
+                    }
                 }
             }
         }