github/reth
1
1
Fork 0
mirror of https://github.com/paradigmxyz/reth.git synced 2026-01-09 23:38:10 -05:00
Code Issues Packages Projects Releases Wiki Activity

chore(trie): Allow reusing Hashed/TrieCursors (#19588)

Browse Source
This commit is contained in:
Brian Picciano
2025-11-12 18:31:04 +01:00
committed by GitHub
parent 95b8a8535b
commit 573191e1d1
18 changed files with 402 additions and 122 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
crates/storage/provider/src/providers/database/provider.rs
View File

@@ -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<TX: DbTxMut + DbTx + 'static, N: NodeTypes> 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<BranchNodeCompact>)> = 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<TX: DbTxMut + DbTx + 'static, N: NodeTypes> 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<BranchNodeCompact>)> = 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<TX: DbTxMut + DbTx + 'static, N: NodeTypes> 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<TX: DbTxMut + DbTx + 'static, N: NodeTypes> 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() {

12
crates/trie/db/src/hashed_cursor.rs
View File

@@ -69,6 +69,10 @@ where
fn next(&mut self) -> Result<Option<(B256, Self::Value)>, 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<Option<(B256, Self::Value)>, 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<C> HashedStorageCursor for DatabaseHashedStorageCursor<C>
@@ -111,4 +119,8 @@ where
fn is_storage_empty(&mut self) -> Result<bool, DatabaseError> {
Ok(self.cursor.seek_exact(self.hashed_address)?.is_none())
}
fn set_hashed_address(&mut self, hashed_address: B256) {
self.hashed_address = hashed_address;
}
}

19
crates/trie/db/src/trie_cursor.rs
View File

@@ -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<Option<Nibbles>, 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<Option<Nibbles>, DatabaseError> {
Ok(self.cursor.current()?.map(|(_, v)| v.nibbles.0))
}
fn reset(&mut self) {
// No-op for database cursors
}
}
impl<C> TrieStorageCursor for DatabaseStorageTrieCursor<C>
where
C: DbCursorRO<tables::StoragesTrie> + DbDupCursorRO<tables::StoragesTrie> + Send + Sync,
{
fn set_hashed_address(&mut self, hashed_address: B256) {
self.hashed_address = hashed_address;
}
}
#[cfg(test)]

6
crates/trie/sparse-parallel/src/trie.rs
View File

@@ -2990,9 +2990,9 @@ mod tests {
.into_sorted();
let mut node_iter = TrieNodeIter::state_trie(
walker,
HashedPostStateCursor::new(
Option::<NoopHashedCursor<Account>>::None,
hashed_post_state.accounts(),
HashedPostStateCursor::new_account(
NoopHashedCursor::<Account>::default(),
&hashed_post_state,
),
);

32
crates/trie/sparse/benches/root.rs
View File

@@ -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::<NoopHashedCursor<U256>>::None,
&storage_sorted.storage_slots,
HashedPostStateCursor::new_storage(
NoopHashedCursor::<U256>::default(),
&hashed_post_state,
hashed_address,
),
);

6
crates/trie/sparse/src/trie.rs
View File

@@ -2416,9 +2416,9 @@ mod tests {
.into_sorted();
let mut node_iter = TrieNodeIter::state_trie(
walker,
HashedPostStateCursor::new(
Option::<NoopHashedCursor<Account>>::None,
hashed_post_state.accounts(),
HashedPostStateCursor::new_account(
NoopHashedCursor::<Account>::default(),
&hashed_post_state,
),
);

6
crates/trie/trie/src/forward_cursor.rs
View File

@@ -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)?;

8
crates/trie/trie/src/hashed_cursor/metrics.rs
View File

@@ -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)
}
}

8
crates/trie/trie/src/hashed_cursor/mock.rs
View File

@@ -145,6 +145,10 @@ impl<T: Debug + Clone> HashedCursor for MockHashedCursor<T> {
});
Ok(entry)
}
fn reset(&mut self) {
self.current_key = None;
}
}
impl<T: Debug + Clone> HashedStorageCursor for MockHashedCursor<T> {
@@ -152,4 +156,8 @@ impl<T: Debug + Clone> HashedStorageCursor for MockHashedCursor<T> {
fn is_storage_empty(&mut self) -> Result<bool, DatabaseError> {
Ok(self.values.is_empty())
}
fn set_hashed_address(&mut self, _hashed_address: B256) {
unimplemented!()
}
}

14
crates/trie/trie/src/hashed_cursor/mod.rs
View File

@@ -53,6 +53,13 @@ pub trait HashedCursor {
/// Move the cursor to the next entry and return it.
fn next(&mut self) -> Result<Option<(B256, Self::Value)>, 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<bool, DatabaseError>;
/// 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);
}

8
crates/trie/trie/src/hashed_cursor/noop.rs
View File

@@ -56,10 +56,18 @@ where
fn next(&mut self) -> Result<Option<(B256, Self::Value)>, DatabaseError> {
Ok(None)
}
fn reset(&mut self) {
// Noop
}
}
impl HashedStorageCursor for NoopHashedCursor<U256> {
fn is_storage_empty(&mut self) -> Result<bool, DatabaseError> {
Ok(true)
}
fn set_hashed_address(&mut self, _hashed_address: B256) {
// Noop
}
}

155
crates/trie/trie/src/hashed_cursor/post_state.rs
View File

@@ -35,27 +35,16 @@ where
fn hashed_account_cursor(&self) -> Result<Self::AccountCursor<'_>, 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<Self::StorageCursor<'_>, 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<C>,
/// 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<Account>>
where
C: HashedCursor<Value = Account>,
{
/// 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<Value = U256>,
{
/// 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<Value = V::NonZero>,
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<C>, 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<C, V> HashedStorageCursor for HashedPostStateCursor<'_, C, V>
impl<C> HashedStorageCursor for HashedPostStateCursor<'_, C, U256>
where
C: HashedStorageCursor<Value = V::NonZero>,
V: HashedPostStateCursorValue,
C: HashedStorageCursor<Value = U256>,
{
/// 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::<C, U256>::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();

6
crates/trie/trie/src/node_iter.rs
View File

@@ -352,9 +352,9 @@ mod tests {
let mut node_iter = TrieNodeIter::state_trie(
walker,
HashedPostStateCursor::new(
Option::<NoopHashedCursor<Account>>::None,
hashed_post_state.accounts(),
HashedPostStateCursor::new_account(
NoopHashedCursor::<Account>::default(),
&hashed_post_state,
),
);

140
crates/trie/trie/src/trie_cursor/in_memory.rs
View File

@@ -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<Self::AccountTrieCursor<'_>, 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<Self::StorageTrieCursor<'_>, DatabaseError> {
// if the storage trie has no updates then we use this as the in-memory overlay.
const EMPTY_UPDATES: &[(Nibbles, Option<BranchNodeCompact>)] = &[];
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<C>,
/// 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<Nibbles>,
/// 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<C>,
trie_updates: &'a [(Nibbles, Option<BranchNodeCompact>)],
/// 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<BranchNodeCompact>>, 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<C: TrieCursor> TrieCursor for InMemoryTrieCursor<'_, C> {
fn current(&mut self) -> Result<Option<Nibbles>, 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<C: TrieStorageCursor> 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();

13
crates/trie/trie/src/trie_cursor/metrics.rs
View File

@@ -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<Option<Nibbles>, 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)
}
}

12
crates/trie/trie/src/trie_cursor/mock.rs
View File

@@ -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<Option<Nibbles>, 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!()
}
}

22
crates/trie/trie/src/trie_cursor/mod.rs
View File

@@ -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<Option<Nibbles>, 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

16
crates/trie/trie/src/trie_cursor/noop.rs
View File

@@ -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<Option<Nibbles>, 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<Option<Nibbles>, DatabaseError> {
Ok(None)
}
fn reset(&mut self) {
// Noop
}
}
impl TrieStorageCursor for NoopStorageTrieCursor {
fn set_hashed_address(&mut self, _hashed_address: B256) {
// Noop
}
}