diff --git a/crates/storage/provider/src/providers/database/provider.rs b/crates/storage/provider/src/providers/database/provider.rs index 62bde757ee..e83e5546f2 100644 --- a/crates/storage/provider/src/providers/database/provider.rs +++ b/crates/storage/provider/src/providers/database/provider.rs @@ -69,8 +69,7 @@ use reth_trie::{ TrieCursorIter, }, updates::{StorageTrieUpdatesSorted, TrieUpdatesSorted}, - BranchNodeCompact, HashedPostStateSorted, Nibbles, StoredNibbles, StoredNibblesSubKey, - TrieChangeSetsEntry, + HashedPostStateSorted, StoredNibbles, StoredNibblesSubKey, TrieChangeSetsEntry, }; use reth_trie_db::{ DatabaseAccountTrieCursor, DatabaseStorageTrieCursor, DatabaseTrieCursorFactory, @@ -2135,17 +2134,13 @@ impl TrieWriter for DatabaseProvider // Wrap the cursor in DatabaseAccountTrieCursor let mut db_account_cursor = DatabaseAccountTrieCursor::new(curr_values_cursor); - // Static empty array for when updates_overlay is None - static EMPTY_ACCOUNT_UPDATES: Vec<(Nibbles, Option)> = Vec::new(); - - // Get the overlay updates for account trie, or use an empty array - let account_overlay_updates = updates_overlay - .map(|overlay| overlay.account_nodes_ref()) - .unwrap_or(&EMPTY_ACCOUNT_UPDATES); + // Create empty TrieUpdatesSorted for when updates_overlay is None + let empty_updates = TrieUpdatesSorted::default(); + let overlay = updates_overlay.unwrap_or(&empty_updates); // Wrap the cursor in InMemoryTrieCursor with the overlay let mut in_memory_account_cursor = - InMemoryTrieCursor::new(Some(&mut db_account_cursor), account_overlay_updates); + InMemoryTrieCursor::new_account(&mut db_account_cursor, overlay); for (path, _) in trie_updates.account_nodes_ref() { num_entries += 1; @@ -2383,8 +2378,8 @@ impl StorageTrieWriter for DatabaseP B256::default(), // Will be set per iteration ); - // Static empty array for when updates_overlay is None - static EMPTY_UPDATES: Vec<(Nibbles, Option)> = Vec::new(); + // Create empty TrieUpdatesSorted for when updates_overlay is None + let empty_updates = TrieUpdatesSorted::default(); for (hashed_address, storage_trie_updates) in storage_tries { let changeset_key = BlockNumberHashedAddress((block_number, *hashed_address)); @@ -2393,15 +2388,15 @@ impl StorageTrieWriter for DatabaseP changed_curr_values_cursor = DatabaseStorageTrieCursor::new(changed_curr_values_cursor.cursor, *hashed_address); - // Get the overlay updates for this storage trie, or use an empty array - let overlay_updates = updates_overlay - .and_then(|overlay| overlay.storage_tries_ref().get(hashed_address)) - .map(|updates| updates.storage_nodes_ref()) - .unwrap_or(&EMPTY_UPDATES); + // Get the overlay updates, or use empty updates + let overlay = updates_overlay.unwrap_or(&empty_updates); // Wrap the cursor in InMemoryTrieCursor with the overlay - let mut in_memory_changed_cursor = - InMemoryTrieCursor::new(Some(&mut changed_curr_values_cursor), overlay_updates); + let mut in_memory_changed_cursor = InMemoryTrieCursor::new_storage( + &mut changed_curr_values_cursor, + overlay, + *hashed_address, + ); // Create an iterator which produces the current values of all updated paths, or None if // they are currently unset. @@ -2417,8 +2412,11 @@ impl StorageTrieWriter for DatabaseP DatabaseStorageTrieCursor::new(wiped_nodes_cursor.cursor, *hashed_address); // Wrap the wiped nodes cursor in InMemoryTrieCursor with the overlay - let mut in_memory_wiped_cursor = - InMemoryTrieCursor::new(Some(&mut wiped_nodes_cursor), overlay_updates); + let mut in_memory_wiped_cursor = InMemoryTrieCursor::new_storage( + &mut wiped_nodes_cursor, + overlay, + *hashed_address, + ); let all_nodes = TrieCursorIter::new(&mut in_memory_wiped_cursor); @@ -3175,6 +3173,7 @@ mod tests { BlockWriter, }; use reth_testing_utils::generators::{self, random_block, BlockParams}; + use reth_trie::Nibbles; #[test] fn test_receipts_by_block_range_empty_range() { diff --git a/crates/trie/db/src/hashed_cursor.rs b/crates/trie/db/src/hashed_cursor.rs index 4fe3d57429..10a1fd8363 100644 --- a/crates/trie/db/src/hashed_cursor.rs +++ b/crates/trie/db/src/hashed_cursor.rs @@ -69,6 +69,10 @@ where fn next(&mut self) -> Result, DatabaseError> { self.0.next() } + + fn reset(&mut self) { + // Database cursors are stateless, no reset needed + } } /// The structure wrapping a database cursor for hashed storage and @@ -102,6 +106,10 @@ where fn next(&mut self) -> Result, DatabaseError> { Ok(self.cursor.next_dup_val()?.map(|e| (e.key, e.value))) } + + fn reset(&mut self) { + // Database cursors are stateless, no reset needed + } } impl HashedStorageCursor for DatabaseHashedStorageCursor @@ -111,4 +119,8 @@ where fn is_storage_empty(&mut self) -> Result { Ok(self.cursor.seek_exact(self.hashed_address)?.is_none()) } + + fn set_hashed_address(&mut self, hashed_address: B256) { + self.hashed_address = hashed_address; + } } diff --git a/crates/trie/db/src/trie_cursor.rs b/crates/trie/db/src/trie_cursor.rs index d05c3fd92d..7b9c402545 100644 --- a/crates/trie/db/src/trie_cursor.rs +++ b/crates/trie/db/src/trie_cursor.rs @@ -6,7 +6,7 @@ use reth_db_api::{ DatabaseError, }; use reth_trie::{ - trie_cursor::{TrieCursor, TrieCursorFactory}, + trie_cursor::{TrieCursor, TrieCursorFactory, TrieStorageCursor}, updates::StorageTrieUpdatesSorted, BranchNodeCompact, Nibbles, StorageTrieEntry, StoredNibbles, StoredNibblesSubKey, }; @@ -91,6 +91,10 @@ where fn current(&mut self) -> Result, DatabaseError> { Ok(self.0.current()?.map(|(k, _)| k.0)) } + + fn reset(&mut self) { + // No-op for database cursors + } } /// A cursor over the storage tries stored in the database. @@ -190,6 +194,19 @@ where fn current(&mut self) -> Result, DatabaseError> { Ok(self.cursor.current()?.map(|(_, v)| v.nibbles.0)) } + + fn reset(&mut self) { + // No-op for database cursors + } +} + +impl TrieStorageCursor for DatabaseStorageTrieCursor +where + C: DbCursorRO + DbDupCursorRO + Send + Sync, +{ + fn set_hashed_address(&mut self, hashed_address: B256) { + self.hashed_address = hashed_address; + } } #[cfg(test)] diff --git a/crates/trie/sparse-parallel/src/trie.rs b/crates/trie/sparse-parallel/src/trie.rs index 4244f20ab1..ab0b9e222d 100644 --- a/crates/trie/sparse-parallel/src/trie.rs +++ b/crates/trie/sparse-parallel/src/trie.rs @@ -2990,9 +2990,9 @@ mod tests { .into_sorted(); let mut node_iter = TrieNodeIter::state_trie( walker, - HashedPostStateCursor::new( - Option::>::None, - hashed_post_state.accounts(), + HashedPostStateCursor::new_account( + NoopHashedCursor::::default(), + &hashed_post_state, ), ); diff --git a/crates/trie/sparse/benches/root.rs b/crates/trie/sparse/benches/root.rs index 76eaac91a7..ece0aa5313 100644 --- a/crates/trie/sparse/benches/root.rs +++ b/crates/trie/sparse/benches/root.rs @@ -12,7 +12,7 @@ use reth_trie::{ walker::TrieWalker, HashedStorage, }; -use reth_trie_common::{HashBuilder, Nibbles}; +use reth_trie_common::{updates::TrieUpdatesSorted, HashBuilder, Nibbles}; use reth_trie_sparse::{provider::DefaultTrieNodeProvider, SerialSparseTrie, SparseTrie}; fn calculate_root_from_leaves(c: &mut Criterion) { @@ -133,18 +133,36 @@ fn calculate_root_from_leaves_repeated(c: &mut Criterion) { ) }; + // Create a TrieUpdatesSorted with just this storage trie + let mut storage_tries = Default::default(); + alloy_primitives::map::B256Map::insert( + &mut storage_tries, + B256::ZERO, + trie_updates_sorted.clone(), + ); + let full_trie_updates = + TrieUpdatesSorted::new(Vec::new(), storage_tries); + let walker = TrieWalker::<_>::storage_trie( - InMemoryTrieCursor::new( - Some(NoopStorageTrieCursor::default()), - &trie_updates_sorted.storage_nodes, + InMemoryTrieCursor::new_storage( + NoopStorageTrieCursor::default(), + &full_trie_updates, + B256::ZERO, ), prefix_set, ); + let hashed_address = B256::ZERO; + let mut storages = alloy_primitives::map::B256Map::default(); + storages.insert(hashed_address, storage_sorted.clone()); + let hashed_post_state = + reth_trie::HashedPostStateSorted::new(Vec::new(), storages); + let mut node_iter = TrieNodeIter::storage_trie( walker, - HashedPostStateCursor::new( - Option::>::None, - &storage_sorted.storage_slots, + HashedPostStateCursor::new_storage( + NoopHashedCursor::::default(), + &hashed_post_state, + hashed_address, ), ); diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs index cf92942b82..d115de8016 100644 --- a/crates/trie/sparse/src/trie.rs +++ b/crates/trie/sparse/src/trie.rs @@ -2416,9 +2416,9 @@ mod tests { .into_sorted(); let mut node_iter = TrieNodeIter::state_trie( walker, - HashedPostStateCursor::new( - Option::>::None, - hashed_post_state.accounts(), + HashedPostStateCursor::new_account( + NoopHashedCursor::::default(), + &hashed_post_state, ), ); diff --git a/crates/trie/trie/src/forward_cursor.rs b/crates/trie/trie/src/forward_cursor.rs index 53d07d5247..5abb5e2431 100644 --- a/crates/trie/trie/src/forward_cursor.rs +++ b/crates/trie/trie/src/forward_cursor.rs @@ -39,6 +39,12 @@ impl<'a, K, V> ForwardInMemoryCursor<'a, K, V> { self.entries.get(self.idx) } + /// Resets the cursor to the beginning of the collection. + #[inline] + pub const fn reset(&mut self) { + self.idx = 0; + } + #[inline] fn next(&mut self) -> Option<&(K, V)> { let entry = self.entries.get(self.idx)?; diff --git a/crates/trie/trie/src/hashed_cursor/metrics.rs b/crates/trie/trie/src/hashed_cursor/metrics.rs index c72ef75289..d7961638e2 100644 --- a/crates/trie/trie/src/hashed_cursor/metrics.rs +++ b/crates/trie/trie/src/hashed_cursor/metrics.rs @@ -151,6 +151,10 @@ where self.metrics.total_duration += start.elapsed(); result } + + fn reset(&mut self) { + self.cursor.reset() + } } impl<'metrics, C> HashedStorageCursor for InstrumentedHashedCursor<'metrics, C> @@ -164,4 +168,8 @@ where self.metrics.total_duration += start.elapsed(); result } + + fn set_hashed_address(&mut self, hashed_address: B256) { + self.cursor.set_hashed_address(hashed_address) + } } diff --git a/crates/trie/trie/src/hashed_cursor/mock.rs b/crates/trie/trie/src/hashed_cursor/mock.rs index 334b5ba77c..ce5b830544 100644 --- a/crates/trie/trie/src/hashed_cursor/mock.rs +++ b/crates/trie/trie/src/hashed_cursor/mock.rs @@ -145,6 +145,10 @@ impl HashedCursor for MockHashedCursor { }); Ok(entry) } + + fn reset(&mut self) { + self.current_key = None; + } } impl HashedStorageCursor for MockHashedCursor { @@ -152,4 +156,8 @@ impl HashedStorageCursor for MockHashedCursor { fn is_storage_empty(&mut self) -> Result { Ok(self.values.is_empty()) } + + fn set_hashed_address(&mut self, _hashed_address: B256) { + unimplemented!() + } } diff --git a/crates/trie/trie/src/hashed_cursor/mod.rs b/crates/trie/trie/src/hashed_cursor/mod.rs index b8f4b08598..90d70c1c68 100644 --- a/crates/trie/trie/src/hashed_cursor/mod.rs +++ b/crates/trie/trie/src/hashed_cursor/mod.rs @@ -53,6 +53,13 @@ pub trait HashedCursor { /// Move the cursor to the next entry and return it. fn next(&mut self) -> Result, DatabaseError>; + + /// Reset the cursor to its initial state. + /// + /// # Important + /// + /// After calling this method, the subsequent operation MUST be a [`HashedCursor::seek`] call. + fn reset(&mut self); } /// The cursor for iterating over hashed storage entries. @@ -60,4 +67,11 @@ pub trait HashedCursor { pub trait HashedStorageCursor: HashedCursor { /// Returns `true` if there are no entries for a given key. fn is_storage_empty(&mut self) -> Result; + + /// Set the hashed address for the storage cursor. + /// + /// # Important + /// + /// After calling this method, the subsequent operation MUST be a [`HashedCursor::seek`] call. + fn set_hashed_address(&mut self, hashed_address: B256); } diff --git a/crates/trie/trie/src/hashed_cursor/noop.rs b/crates/trie/trie/src/hashed_cursor/noop.rs index 4f59af2ed3..88726d3b67 100644 --- a/crates/trie/trie/src/hashed_cursor/noop.rs +++ b/crates/trie/trie/src/hashed_cursor/noop.rs @@ -56,10 +56,18 @@ where fn next(&mut self) -> Result, DatabaseError> { Ok(None) } + + fn reset(&mut self) { + // Noop + } } impl HashedStorageCursor for NoopHashedCursor { fn is_storage_empty(&mut self) -> Result { Ok(true) } + + fn set_hashed_address(&mut self, _hashed_address: B256) { + // Noop + } } diff --git a/crates/trie/trie/src/hashed_cursor/post_state.rs b/crates/trie/trie/src/hashed_cursor/post_state.rs index bf80ee63a1..0b78119baf 100644 --- a/crates/trie/trie/src/hashed_cursor/post_state.rs +++ b/crates/trie/trie/src/hashed_cursor/post_state.rs @@ -35,27 +35,16 @@ where fn hashed_account_cursor(&self) -> Result, DatabaseError> { let cursor = self.cursor_factory.hashed_account_cursor()?; - Ok(HashedPostStateCursor::new(Some(cursor), &self.post_state.as_ref().accounts)) + Ok(HashedPostStateCursor::new_account(cursor, self.post_state.as_ref())) } fn hashed_storage_cursor( &self, hashed_address: B256, ) -> Result, DatabaseError> { - static EMPTY_UPDATES: Vec<(B256, U256)> = Vec::new(); - - let post_state_storage = self.post_state.as_ref().storages.get(&hashed_address); - let (storage_slots, wiped) = post_state_storage - .map(|u| (u.storage_slots_ref(), u.is_wiped())) - .unwrap_or((&EMPTY_UPDATES, false)); - - let cursor = if wiped { - None - } else { - Some(self.cursor_factory.hashed_storage_cursor(hashed_address)?) - }; - - Ok(HashedPostStateCursor::new(cursor, storage_slots)) + let post_state = self.post_state.as_ref(); + let cursor = self.cursor_factory.hashed_storage_cursor(hashed_address)?; + Ok(HashedPostStateCursor::new_storage(cursor, post_state, hashed_address)) } } @@ -101,8 +90,10 @@ pub struct HashedPostStateCursor<'a, C, V> where V: HashedPostStateCursorValue, { - /// The underlying `database_cursor`. If None then it is assumed there is no DB data. - cursor: Option, + /// The underlying cursor. + cursor: C, + /// Whether the underlying cursor should be ignored (when storage was wiped). + cursor_wiped: bool, /// Entry that `database_cursor` is currently pointing to. cursor_entry: Option<(B256, V::NonZero)>, /// Forward-only in-memory cursor over underlying V. @@ -113,6 +104,64 @@ where /// Tracks whether `seek` has been called. Used to prevent re-seeking the DB cursor /// when it has been exhausted by iteration. seeked: bool, + /// Reference to the full post state. + post_state: &'a HashedPostStateSorted, +} + +impl<'a, C> HashedPostStateCursor<'a, C, Option> +where + C: HashedCursor, +{ + /// Create new account cursor which combines a DB cursor and the post state. + pub fn new_account(cursor: C, post_state: &'a HashedPostStateSorted) -> Self { + let post_state_cursor = ForwardInMemoryCursor::new(&post_state.accounts); + Self { + cursor, + cursor_wiped: false, + cursor_entry: None, + post_state_cursor, + last_key: None, + seeked: false, + post_state, + } + } +} + +impl<'a, C> HashedPostStateCursor<'a, C, U256> +where + C: HashedStorageCursor, +{ + /// Create new storage cursor with full post state reference. + /// This allows the cursor to switch between storage tries when `set_hashed_address` is called. + pub fn new_storage( + cursor: C, + post_state: &'a HashedPostStateSorted, + hashed_address: B256, + ) -> Self { + let (post_state_cursor, cursor_wiped) = + Self::get_storage_overlay(post_state, hashed_address); + Self { + cursor, + cursor_wiped, + cursor_entry: None, + post_state_cursor, + last_key: None, + seeked: false, + post_state, + } + } + + /// Returns the storage overlay for `hashed_address` and whether it was wiped. + fn get_storage_overlay( + post_state: &'a HashedPostStateSorted, + hashed_address: B256, + ) -> (ForwardInMemoryCursor<'a, B256, U256>, bool) { + let post_state_storage = post_state.storages.get(&hashed_address); + let cursor_wiped = post_state_storage.is_some_and(|u| u.is_wiped()); + let storage_slots = post_state_storage.map(|u| u.storage_slots_ref()).unwrap_or(&[]); + + (ForwardInMemoryCursor::new(storage_slots), cursor_wiped) + } } impl<'a, C, V> HashedPostStateCursor<'a, C, V> @@ -120,22 +169,9 @@ where C: HashedCursor, V: HashedPostStateCursorValue, { - /// Creates a new post state cursor which combines a DB cursor and in-memory post state updates. - /// - /// # Parameters - /// - `cursor`: The database cursor. Pass `None` to indicate: - /// - For accounts: Empty database (no persisted accounts) - /// - For storage: Wiped storage (e.g., via `SELFDESTRUCT` - all previous storage destroyed) - /// - `updates`: Pre-sorted post state updates. - pub fn new(cursor: Option, updates: &'a [(B256, V)]) -> Self { - debug_assert!(updates.is_sorted_by_key(|(k, _)| k), "Overlay values must be sorted by key"); - Self { - cursor, - cursor_entry: None, - post_state_cursor: ForwardInMemoryCursor::new(updates), - last_key: None, - seeked: false, - } + /// Returns a mutable reference to the underlying cursor if it's not wiped, None otherwise. + fn get_cursor_mut(&mut self) -> Option<&mut C> { + (!self.cursor_wiped).then_some(&mut self.cursor) } /// Asserts that the next entry to be returned from the cursor is not previous to the last entry @@ -162,7 +198,7 @@ where }; if should_seek { - self.cursor_entry = self.cursor.as_mut().map(|c| c.seek(key)).transpose()?.flatten(); + self.cursor_entry = self.get_cursor_mut().map(|c| c.seek(key)).transpose()?.flatten(); } Ok(()) @@ -175,7 +211,7 @@ where // If the previous entry is `None`, and we've done a seek previously, then the cursor is // exhausted, and we shouldn't call `next` again. if self.cursor_entry.is_some() { - self.cursor_entry = self.cursor.as_mut().map(|c| c.next()).transpose()?.flatten(); + self.cursor_entry = self.get_cursor_mut().map(|c| c.next()).transpose()?.flatten(); } Ok(()) @@ -281,14 +317,33 @@ where self.set_last_key(&entry); Ok(entry) } + + fn reset(&mut self) { + let Self { + cursor, + cursor_wiped, + cursor_entry, + post_state_cursor, + last_key, + seeked, + post_state: _, + } = self; + + cursor.reset(); + post_state_cursor.reset(); + + *cursor_wiped = false; + *cursor_entry = None; + *last_key = None; + *seeked = false; + } } /// The cursor to iterate over post state hashed values and corresponding database entries. /// It will always give precedence to the data from the post state. -impl HashedStorageCursor for HashedPostStateCursor<'_, C, V> +impl HashedStorageCursor for HashedPostStateCursor<'_, C, U256> where - C: HashedStorageCursor, - V: HashedPostStateCursorValue, + C: HashedStorageCursor, { /// Returns `true` if the account has no storage entries. /// @@ -301,8 +356,15 @@ where } // If no non-zero slots in post state, check the database. - // Returns true if cursor is None (wiped storage or empty DB). - self.cursor.as_mut().map_or(Ok(true), |c| c.is_storage_empty()) + // Returns true if cursor is wiped. + self.get_cursor_mut().map_or(Ok(true), |c| c.is_storage_empty()) + } + + fn set_hashed_address(&mut self, hashed_address: B256) { + self.reset(); + self.cursor.set_hashed_address(hashed_address); + (self.post_state_cursor, self.cursor_wiped) = + HashedPostStateCursor::::get_storage_overlay(self.post_state, hashed_address); } } @@ -417,7 +479,18 @@ mod tests { let db_nodes_arc = Arc::new(db_nodes_map); let visited_keys = Arc::new(Mutex::new(Vec::new())); let mock_cursor = MockHashedCursor::new(db_nodes_arc, visited_keys); - let mut test_cursor = HashedPostStateCursor::new(Some(mock_cursor), &post_state_nodes); + + // Create a HashedPostStateSorted with the storage data + let hashed_address = B256::ZERO; + let storage_sorted = reth_trie_common::HashedStorageSorted { + storage_slots: post_state_nodes, + wiped: false, + }; + let mut storages = alloy_primitives::map::B256Map::default(); + storages.insert(hashed_address, storage_sorted); + let post_state = HashedPostStateSorted::new(Vec::new(), storages); + + let mut test_cursor = HashedPostStateCursor::new_storage(mock_cursor, &post_state, hashed_address); // Test: seek to the beginning first let control_first = control_cursor.seek(B256::ZERO).unwrap(); diff --git a/crates/trie/trie/src/node_iter.rs b/crates/trie/trie/src/node_iter.rs index aef668e8ff..b57fc2da70 100644 --- a/crates/trie/trie/src/node_iter.rs +++ b/crates/trie/trie/src/node_iter.rs @@ -352,9 +352,9 @@ mod tests { let mut node_iter = TrieNodeIter::state_trie( walker, - HashedPostStateCursor::new( - Option::>::None, - hashed_post_state.accounts(), + HashedPostStateCursor::new_account( + NoopHashedCursor::::default(), + &hashed_post_state, ), ); diff --git a/crates/trie/trie/src/trie_cursor/in_memory.rs b/crates/trie/trie/src/trie_cursor/in_memory.rs index 308a8edbd0..941fbf9633 100644 --- a/crates/trie/trie/src/trie_cursor/in_memory.rs +++ b/crates/trie/trie/src/trie_cursor/in_memory.rs @@ -1,4 +1,4 @@ -use super::{TrieCursor, TrieCursorFactory}; +use super::{TrieCursor, TrieCursorFactory, TrieStorageCursor}; use crate::{forward_cursor::ForwardInMemoryCursor, updates::TrieUpdatesSorted}; use alloy_primitives::B256; use reth_storage_errors::db::DatabaseError; @@ -37,29 +37,16 @@ where fn account_trie_cursor(&self) -> Result, DatabaseError> { let cursor = self.cursor_factory.account_trie_cursor()?; - Ok(InMemoryTrieCursor::new(Some(cursor), self.trie_updates.as_ref().account_nodes_ref())) + Ok(InMemoryTrieCursor::new_account(cursor, self.trie_updates.as_ref())) } fn storage_trie_cursor( &self, hashed_address: B256, ) -> Result, DatabaseError> { - // if the storage trie has no updates then we use this as the in-memory overlay. - const EMPTY_UPDATES: &[(Nibbles, Option)] = &[]; - - let storage_trie_updates = - self.trie_updates.as_ref().storage_tries_ref().get(&hashed_address); - let (storage_nodes, cleared) = storage_trie_updates - .map(|u| (u.storage_nodes_ref(), u.is_deleted())) - .unwrap_or((EMPTY_UPDATES, false)); - - let cursor = if cleared { - None - } else { - Some(self.cursor_factory.storage_trie_cursor(hashed_address)?) - }; - - Ok(InMemoryTrieCursor::new(cursor, storage_nodes)) + let trie_updates = self.trie_updates.as_ref(); + let cursor = self.cursor_factory.storage_trie_cursor(hashed_address)?; + Ok(InMemoryTrieCursor::new_storage(cursor, trie_updates, hashed_address)) } } @@ -67,8 +54,10 @@ where /// It will always give precedence to the data from the trie updates. #[derive(Debug)] pub struct InMemoryTrieCursor<'a, C> { - /// The underlying cursor. If None then it is assumed there is no DB data. - cursor: Option, + /// The underlying cursor. + cursor: C, + /// Whether the underlying cursor should be ignored (when storage trie was wiped). + cursor_wiped: bool, /// Entry that `cursor` is currently pointing to. cursor_entry: Option<(Nibbles, BranchNodeCompact)>, /// Forward-only in-memory cursor over storage trie nodes. @@ -77,21 +66,60 @@ pub struct InMemoryTrieCursor<'a, C> { last_key: Option, /// Whether an initial seek was called. seeked: bool, + /// Reference to the full trie updates. + trie_updates: &'a TrieUpdatesSorted, } impl<'a, C: TrieCursor> InMemoryTrieCursor<'a, C> { - /// Create new trie cursor which combines a DB cursor (None to assume empty DB) and a set of - /// in-memory trie nodes. - pub fn new( - cursor: Option, - trie_updates: &'a [(Nibbles, Option)], + /// Create new account trie cursor which combines a DB cursor and the trie updates. + pub fn new_account(cursor: C, trie_updates: &'a TrieUpdatesSorted) -> Self { + let in_memory_cursor = ForwardInMemoryCursor::new(trie_updates.account_nodes_ref()); + Self { + cursor, + cursor_wiped: false, + cursor_entry: None, + in_memory_cursor, + last_key: None, + seeked: false, + trie_updates, + } + } + + /// Create new storage trie cursor with full trie updates reference. + /// This allows the cursor to switch between storage tries when `set_hashed_address` is called. + pub fn new_storage( + cursor: C, + trie_updates: &'a TrieUpdatesSorted, + hashed_address: B256, ) -> Self { - debug_assert!( - trie_updates.is_sorted_by_key(|(k, _)| k), - "Overlay values must be sorted by path" - ); - let in_memory_cursor = ForwardInMemoryCursor::new(trie_updates); - Self { cursor, cursor_entry: None, in_memory_cursor, last_key: None, seeked: false } + let (in_memory_cursor, cursor_wiped) = + Self::get_storage_overlay(trie_updates, hashed_address); + Self { + cursor, + cursor_wiped, + cursor_entry: None, + in_memory_cursor, + last_key: None, + seeked: false, + trie_updates, + } + } + + /// Returns the storage overlay for `hashed_address` and whether it was deleted. + fn get_storage_overlay( + trie_updates: &'a TrieUpdatesSorted, + hashed_address: B256, + ) -> (ForwardInMemoryCursor<'a, Nibbles, Option>, bool) { + let storage_trie_updates = trie_updates.storage_tries_ref().get(&hashed_address); + let cursor_wiped = storage_trie_updates.is_some_and(|u| u.is_deleted()); + let storage_nodes = storage_trie_updates.map(|u| u.storage_nodes_ref()).unwrap_or(&[]); + + (ForwardInMemoryCursor::new(storage_nodes), cursor_wiped) + } + + /// Returns a mutable reference to the underlying cursor if it's not wiped, None otherwise. + fn get_cursor_mut(&mut self) -> Option<&mut C> { + (!self.cursor_wiped).then_some(&mut self.cursor) } /// Asserts that the next entry to be returned from the cursor is not previous to the last entry @@ -118,7 +146,7 @@ impl<'a, C: TrieCursor> InMemoryTrieCursor<'a, C> { }; if should_seek { - self.cursor_entry = self.cursor.as_mut().map(|c| c.seek(key)).transpose()?.flatten(); + self.cursor_entry = self.get_cursor_mut().map(|c| c.seek(key)).transpose()?.flatten(); } Ok(()) @@ -131,7 +159,7 @@ impl<'a, C: TrieCursor> InMemoryTrieCursor<'a, C> { // If the previous entry is `None`, and we've done a seek previously, then the cursor is // exhausted and we shouldn't call `next` again. if self.cursor_entry.is_some() { - self.cursor_entry = self.cursor.as_mut().map(|c| c.next()).transpose()?.flatten(); + self.cursor_entry = self.get_cursor_mut().map(|c| c.next()).transpose()?.flatten(); } Ok(()) @@ -242,9 +270,38 @@ impl TrieCursor for InMemoryTrieCursor<'_, C> { fn current(&mut self) -> Result, DatabaseError> { match &self.last_key { Some(key) => Ok(Some(*key)), - None => Ok(self.cursor.as_mut().map(|c| c.current()).transpose()?.flatten()), + None => Ok(self.get_cursor_mut().map(|c| c.current()).transpose()?.flatten()), } } + + fn reset(&mut self) { + let Self { + cursor, + cursor_wiped, + cursor_entry, + in_memory_cursor, + last_key, + seeked, + trie_updates: _, + } = self; + + cursor.reset(); + in_memory_cursor.reset(); + + *cursor_wiped = false; + *cursor_entry = None; + *last_key = None; + *seeked = false; + } +} + +impl TrieStorageCursor for InMemoryTrieCursor<'_, C> { + fn set_hashed_address(&mut self, hashed_address: B256) { + self.reset(); + self.cursor.set_hashed_address(hashed_address); + (self.in_memory_cursor, self.cursor_wiped) = + Self::get_storage_overlay(self.trie_updates, hashed_address); + } } #[cfg(test)] @@ -268,7 +325,8 @@ mod tests { let visited_keys = Arc::new(Mutex::new(Vec::new())); let mock_cursor = MockTrieCursor::new(db_nodes_arc, visited_keys); - let mut cursor = InMemoryTrieCursor::new(Some(mock_cursor), &test_case.in_memory_nodes); + let trie_updates = TrieUpdatesSorted::new(test_case.in_memory_nodes, Default::default()); + let mut cursor = InMemoryTrieCursor::new_account(mock_cursor, &trie_updates); let mut results = Vec::new(); @@ -451,7 +509,8 @@ mod tests { let visited_keys = Arc::new(Mutex::new(Vec::new())); let mock_cursor = MockTrieCursor::new(db_nodes_arc, visited_keys); - let mut cursor = InMemoryTrieCursor::new(Some(mock_cursor), &in_memory_nodes); + let trie_updates = TrieUpdatesSorted::new(in_memory_nodes, Default::default()); + let mut cursor = InMemoryTrieCursor::new_account(mock_cursor, &trie_updates); let result = cursor.seek_exact(Nibbles::from_nibbles([0x2])).unwrap(); assert_eq!( @@ -550,7 +609,8 @@ mod tests { let visited_keys = Arc::new(Mutex::new(Vec::new())); let mock_cursor = MockTrieCursor::new(db_nodes_arc, visited_keys); - let mut cursor = InMemoryTrieCursor::new(Some(mock_cursor), &in_memory_nodes); + let trie_updates = TrieUpdatesSorted::new(in_memory_nodes, Default::default()); + let mut cursor = InMemoryTrieCursor::new_account(mock_cursor, &trie_updates); assert_eq!(cursor.current().unwrap(), None); @@ -600,7 +660,8 @@ mod tests { let visited_keys = Arc::new(Mutex::new(Vec::new())); let mock_cursor = MockTrieCursor::new(db_nodes_arc, visited_keys); - let mut cursor = InMemoryTrieCursor::new(Some(mock_cursor), &in_memory_nodes); + let trie_updates = TrieUpdatesSorted::new(in_memory_nodes, Default::default()); + let mut cursor = InMemoryTrieCursor::new_account(mock_cursor, &trie_updates); // Seek to beginning should return None (all nodes are deleted) tracing::debug!("seeking to 0x"); @@ -764,7 +825,8 @@ mod tests { let db_nodes_arc = Arc::new(db_nodes_map); let visited_keys = Arc::new(Mutex::new(Vec::new())); let mock_cursor = MockTrieCursor::new(db_nodes_arc, visited_keys); - let mut test_cursor = InMemoryTrieCursor::new(Some(mock_cursor), &in_memory_nodes); + let trie_updates = TrieUpdatesSorted::new(in_memory_nodes, Default::default()); + let mut test_cursor = InMemoryTrieCursor::new_account(mock_cursor, &trie_updates); // Test: seek to the beginning first let control_first = control_cursor.seek(Nibbles::default()).unwrap(); diff --git a/crates/trie/trie/src/trie_cursor/metrics.rs b/crates/trie/trie/src/trie_cursor/metrics.rs index 8aca09f72b..01d4df1cee 100644 --- a/crates/trie/trie/src/trie_cursor/metrics.rs +++ b/crates/trie/trie/src/trie_cursor/metrics.rs @@ -1,5 +1,6 @@ -use super::TrieCursor; +use super::{TrieCursor, TrieStorageCursor}; use crate::{BranchNodeCompact, Nibbles}; +use alloy_primitives::B256; use reth_storage_errors::db::DatabaseError; use std::time::{Duration, Instant}; @@ -160,4 +161,14 @@ impl<'metrics, C: TrieCursor> TrieCursor for InstrumentedTrieCursor<'metrics, C> fn current(&mut self) -> Result, DatabaseError> { self.cursor.current() } + + fn reset(&mut self) { + self.cursor.reset() + } +} + +impl<'metrics, C: TrieStorageCursor> TrieStorageCursor for InstrumentedTrieCursor<'metrics, C> { + fn set_hashed_address(&mut self, hashed_address: B256) { + self.cursor.set_hashed_address(hashed_address) + } } diff --git a/crates/trie/trie/src/trie_cursor/mock.rs b/crates/trie/trie/src/trie_cursor/mock.rs index 313df0443e..81b5b511f9 100644 --- a/crates/trie/trie/src/trie_cursor/mock.rs +++ b/crates/trie/trie/src/trie_cursor/mock.rs @@ -2,7 +2,7 @@ use parking_lot::{Mutex, MutexGuard}; use std::{collections::BTreeMap, sync::Arc}; use tracing::instrument; -use super::{TrieCursor, TrieCursorFactory}; +use super::{TrieCursor, TrieCursorFactory, TrieStorageCursor}; use crate::{ mock::{KeyVisit, KeyVisitType}, BranchNodeCompact, Nibbles, @@ -165,4 +165,14 @@ impl TrieCursor for MockTrieCursor { fn current(&mut self) -> Result, DatabaseError> { Ok(self.current_key) } + + fn reset(&mut self) { + self.current_key = None; + } +} + +impl TrieStorageCursor for MockTrieCursor { + fn set_hashed_address(&mut self, _hashed_address: B256) { + unimplemented!() + } } diff --git a/crates/trie/trie/src/trie_cursor/mod.rs b/crates/trie/trie/src/trie_cursor/mod.rs index ace57a257a..32fd17c996 100644 --- a/crates/trie/trie/src/trie_cursor/mod.rs +++ b/crates/trie/trie/src/trie_cursor/mod.rs @@ -35,7 +35,7 @@ pub trait TrieCursorFactory { Self: 'a; /// The storage trie cursor type. - type StorageTrieCursor<'a>: TrieCursor + type StorageTrieCursor<'a>: TrieStorageCursor where Self: 'a; @@ -68,6 +68,26 @@ pub trait TrieCursor { /// Get the current entry. fn current(&mut self) -> Result, DatabaseError>; + + /// Reset the cursor to the beginning. + /// + /// # Important + /// + /// After calling this method, the subsequent operation MUST be a [`TrieCursor::seek`] or + /// [`TrieCursor::seek_exact`] call. + fn reset(&mut self); +} + +/// A cursor for traversing storage trie nodes. +#[auto_impl::auto_impl(&mut)] +pub trait TrieStorageCursor: TrieCursor { + /// Set the hashed address for the storage trie cursor. + /// + /// # Important + /// + /// After calling this method, the subsequent operation MUST be a [`TrieCursor::seek`] or + /// [`TrieCursor::seek_exact`] call. + fn set_hashed_address(&mut self, hashed_address: B256); } /// Iterator wrapper for `TrieCursor` types diff --git a/crates/trie/trie/src/trie_cursor/noop.rs b/crates/trie/trie/src/trie_cursor/noop.rs index a00a18e4f0..0a51fd806d 100644 --- a/crates/trie/trie/src/trie_cursor/noop.rs +++ b/crates/trie/trie/src/trie_cursor/noop.rs @@ -1,4 +1,4 @@ -use super::{TrieCursor, TrieCursorFactory}; +use super::{TrieCursor, TrieCursorFactory, TrieStorageCursor}; use crate::{BranchNodeCompact, Nibbles}; use alloy_primitives::B256; use reth_storage_errors::db::DatabaseError; @@ -60,6 +60,10 @@ impl TrieCursor for NoopAccountTrieCursor { fn current(&mut self) -> Result, DatabaseError> { Ok(None) } + + fn reset(&mut self) { + // Noop + } } /// Noop storage trie cursor. @@ -89,4 +93,14 @@ impl TrieCursor for NoopStorageTrieCursor { fn current(&mut self) -> Result, DatabaseError> { Ok(None) } + + fn reset(&mut self) { + // Noop + } +} + +impl TrieStorageCursor for NoopStorageTrieCursor { + fn set_hashed_address(&mut self, _hashed_address: B256) { + // Noop + } }