chore(trie): simplify hashed cursor abstraction (#8380)

This commit is contained in:
Roman Krasiuk
2024-05-27 14:38:39 +02:00
committed by GitHub
parent 2bf5c93039
commit b4d7d368a4
4 changed files with 112 additions and 89 deletions

View File

@ -1,10 +1,10 @@
use super::{HashedAccountCursor, HashedCursorFactory, HashedStorageCursor};
use super::{HashedCursor, HashedCursorFactory, HashedStorageCursor};
use reth_db::{
cursor::{DbCursorRO, DbDupCursorRO},
tables,
transaction::DbTx,
};
use reth_primitives::{Account, StorageEntry, B256};
use reth_primitives::{Account, B256, U256};
impl<'a, TX: DbTx> HashedCursorFactory for &'a TX {
type AccountCursor = <TX as DbTx>::Cursor<tables::HashedAccounts>;
@ -26,25 +26,29 @@ impl<'a, TX: DbTx> HashedCursorFactory for &'a TX {
}
}
impl<C> HashedAccountCursor for C
impl<C> HashedCursor for C
where
C: DbCursorRO<tables::HashedAccounts>,
{
fn seek(&mut self, key: B256) -> Result<Option<(B256, Account)>, reth_db::DatabaseError> {
type Value = Account;
fn seek(&mut self, key: B256) -> Result<Option<(B256, Self::Value)>, reth_db::DatabaseError> {
self.seek(key)
}
fn next(&mut self) -> Result<Option<(B256, Account)>, reth_db::DatabaseError> {
fn next(&mut self) -> Result<Option<(B256, Self::Value)>, reth_db::DatabaseError> {
self.next()
}
}
/// The structure wrapping a database cursor for hashed storage and
/// a target hashed address. Implements [HashedStorageCursor] for iterating
/// hashed state
/// a target hashed address. Implements [HashedCursor] and [HashedStorageCursor]
/// for iterating over hashed storage.
#[derive(Debug)]
pub struct DatabaseHashedStorageCursor<C> {
/// Database hashed storage cursor.
cursor: C,
/// Target hashed address of the account that the storage belongs to.
hashed_address: B256,
}
@ -55,6 +59,24 @@ impl<C> DatabaseHashedStorageCursor<C> {
}
}
impl<C> HashedCursor for DatabaseHashedStorageCursor<C>
where
C: DbCursorRO<tables::HashedStorages> + DbDupCursorRO<tables::HashedStorages>,
{
type Value = U256;
fn seek(
&mut self,
subkey: B256,
) -> Result<Option<(B256, Self::Value)>, reth_db::DatabaseError> {
Ok(self.cursor.seek_by_key_subkey(self.hashed_address, subkey)?.map(|e| (e.key, e.value)))
}
fn next(&mut self) -> Result<Option<(B256, Self::Value)>, reth_db::DatabaseError> {
Ok(self.cursor.next_dup_val()?.map(|e| (e.key, e.value)))
}
}
impl<C> HashedStorageCursor for DatabaseHashedStorageCursor<C>
where
C: DbCursorRO<tables::HashedStorages> + DbDupCursorRO<tables::HashedStorages>,
@ -62,12 +84,4 @@ where
fn is_storage_empty(&mut self) -> Result<bool, reth_db::DatabaseError> {
Ok(self.cursor.seek_exact(self.hashed_address)?.is_none())
}
fn seek(&mut self, subkey: B256) -> Result<Option<StorageEntry>, reth_db::DatabaseError> {
self.cursor.seek_by_key_subkey(self.hashed_address, subkey)
}
fn next(&mut self) -> Result<Option<StorageEntry>, reth_db::DatabaseError> {
self.cursor.next_dup_val()
}
}

View File

@ -1,4 +1,4 @@
use reth_primitives::{Account, StorageEntry, B256};
use reth_primitives::{Account, B256, U256};
/// Default implementation of the hashed state cursor traits.
mod default;
@ -11,9 +11,9 @@ pub use post_state::*;
/// The factory trait for creating cursors over the hashed state.
pub trait HashedCursorFactory {
/// The hashed account cursor type.
type AccountCursor: HashedAccountCursor;
type AccountCursor: HashedCursor<Value = Account>;
/// The hashed storage cursor type.
type StorageCursor: HashedStorageCursor;
type StorageCursor: HashedStorageCursor<Value = U256>;
/// Returns a cursor for iterating over all hashed accounts in the state.
fn hashed_account_cursor(&self) -> Result<Self::AccountCursor, reth_db::DatabaseError>;
@ -25,23 +25,21 @@ pub trait HashedCursorFactory {
) -> Result<Self::StorageCursor, reth_db::DatabaseError>;
}
/// The cursor for iterating over hashed accounts.
pub trait HashedAccountCursor {
/// The cursor for iterating over hashed entries.
pub trait HashedCursor {
/// Value returned by the cursor.
type Value;
/// Seek an entry greater or equal to the given key and position the cursor there.
fn seek(&mut self, key: B256) -> Result<Option<(B256, Account)>, reth_db::DatabaseError>;
/// Returns the first entry with the key greater or equal to the sought key.
fn seek(&mut self, key: B256) -> Result<Option<(B256, Self::Value)>, reth_db::DatabaseError>;
/// Move the cursor to the next entry and return it.
fn next(&mut self) -> Result<Option<(B256, Account)>, reth_db::DatabaseError>;
fn next(&mut self) -> Result<Option<(B256, Self::Value)>, reth_db::DatabaseError>;
}
/// The cursor for iterating over hashed storage entries.
pub trait HashedStorageCursor {
pub trait HashedStorageCursor: HashedCursor {
/// Returns `true` if there are no entries for a given key.
fn is_storage_empty(&mut self) -> Result<bool, reth_db::DatabaseError>;
/// Seek an entry greater or equal to the given key/subkey and position the cursor there.
fn seek(&mut self, subkey: B256) -> Result<Option<StorageEntry>, reth_db::DatabaseError>;
/// Move the cursor to the next entry and return it.
fn next(&mut self) -> Result<Option<StorageEntry>, reth_db::DatabaseError>;
}

View File

@ -1,6 +1,6 @@
use super::{HashedAccountCursor, HashedCursorFactory, HashedStorageCursor};
use super::{HashedCursor, HashedCursorFactory, HashedStorageCursor};
use crate::state::HashedPostStateSorted;
use reth_primitives::{Account, StorageEntry, B256, U256};
use reth_primitives::{Account, B256, U256};
/// The hashed cursor factory for the post state.
#[derive(Debug, Clone)]
@ -88,10 +88,12 @@ impl<'b, C> HashedPostStateAccountCursor<'b, C> {
}
}
impl<'b, C> HashedAccountCursor for HashedPostStateAccountCursor<'b, C>
impl<'b, C> HashedCursor for HashedPostStateAccountCursor<'b, C>
where
C: HashedAccountCursor,
C: HashedCursor<Value = Account>,
{
type Value = Account;
/// Seek the next entry for a given hashed account key.
///
/// If the post state contains the exact match for the key, return it.
@ -99,8 +101,8 @@ where
/// database and the post state. The two entries are compared and the lowest is returned.
///
/// The returned account key is memoized and the cursor remains positioned at that key until
/// [HashedAccountCursor::seek] or [HashedAccountCursor::next] are called.
fn seek(&mut self, key: B256) -> Result<Option<(B256, Account)>, reth_db::DatabaseError> {
/// [HashedCursor::seek] or [HashedCursor::next] are called.
fn seek(&mut self, key: B256) -> Result<Option<(B256, Self::Value)>, reth_db::DatabaseError> {
self.last_account = None;
// Take the next account from the post state with the key greater than or equal to the
@ -142,9 +144,9 @@ where
/// If the cursor is positioned at the entry, return the entry with next greater key.
/// Returns [None] if the previous memoized or the next greater entries are missing.
///
/// NOTE: This function will not return any entry unless [HashedAccountCursor::seek] has been
/// NOTE: This function will not return any entry unless [HashedCursor::seek] has been
/// called.
fn next(&mut self) -> Result<Option<(B256, Account)>, reth_db::DatabaseError> {
fn next(&mut self) -> Result<Option<(B256, Self::Value)>, reth_db::DatabaseError> {
let last_account = match self.last_account.as_ref() {
Some(account) => account,
None => return Ok(None), // no previous entry was found
@ -192,7 +194,7 @@ pub struct HashedPostStateStorageCursor<'b, C> {
}
impl<'b, C> HashedPostStateStorageCursor<'b, C> {
/// Create new instance of [HashedPostStateStorageCursor].
/// Create new instance of [HashedPostStateStorageCursor] for the given hashed address.
pub fn new(cursor: C, post_state: &'b HashedPostStateSorted, hashed_address: B256) -> Self {
Self { cursor, post_state, hashed_address, last_slot: None, post_state_storage_index: 0 }
}
@ -222,49 +224,35 @@ impl<'b, C> HashedPostStateStorageCursor<'b, C> {
/// If the storage keys are the same, the post state entry is given precedence.
fn next_slot(
post_state_item: Option<&(B256, U256)>,
db_item: Option<StorageEntry>,
) -> Option<StorageEntry> {
db_item: Option<(B256, U256)>,
) -> Option<(B256, U256)> {
match (post_state_item, db_item) {
// If both are not empty, return the smallest of the two
// Post state is given precedence if keys are equal
(Some((post_state_slot, post_state_value)), Some(db_entry)) => {
if post_state_slot <= &db_entry.key {
Some(StorageEntry { key: *post_state_slot, value: *post_state_value })
(Some((post_state_slot, post_state_value)), Some((db_slot, db_value))) => {
if post_state_slot <= &db_slot {
Some((*post_state_slot, *post_state_value))
} else {
Some(db_entry)
Some((db_slot, db_value))
}
}
// Return either non-empty entry
_ => db_item.or_else(|| {
post_state_item.copied().map(|(key, value)| StorageEntry { key, value })
}),
_ => db_item.or_else(|| post_state_item.copied()),
}
}
}
impl<'b, C> HashedStorageCursor for HashedPostStateStorageCursor<'b, C>
impl<'b, C> HashedCursor for HashedPostStateStorageCursor<'b, C>
where
C: HashedStorageCursor,
C: HashedStorageCursor<Value = U256>,
{
/// Returns `true` if the account has no storage entries.
///
/// This function should be called before attempting to call [HashedStorageCursor::seek] or
/// [HashedStorageCursor::next].
fn is_storage_empty(&mut self) -> Result<bool, reth_db::DatabaseError> {
let is_empty = match self.post_state.storages.get(&self.hashed_address) {
Some(storage) => {
// If the storage has been wiped at any point
storage.wiped &&
// and the current storage does not contain any non-zero values
storage.non_zero_valued_slots.is_empty()
}
None => self.cursor.is_storage_empty()?,
};
Ok(is_empty)
}
type Value = U256;
/// Seek the next account storage entry for a given hashed key pair.
fn seek(&mut self, subkey: B256) -> Result<Option<StorageEntry>, reth_db::DatabaseError> {
fn seek(
&mut self,
subkey: B256,
) -> Result<Option<(B256, Self::Value)>, reth_db::DatabaseError> {
// Attempt to find the account's storage in post state.
let mut post_state_entry = None;
if let Some(storage) = self.post_state.storages.get(&self.hashed_address) {
@ -281,7 +269,7 @@ where
if let Some((slot, value)) = post_state_entry {
if slot == &subkey {
self.last_slot = Some(*slot);
return Ok(Some(StorageEntry { key: *slot, value: *value }))
return Ok(Some((*slot, *value)))
}
}
@ -293,7 +281,7 @@ where
while db_entry
.as_ref()
.map(|entry| self.is_slot_zero_valued(&entry.key))
.map(|entry| self.is_slot_zero_valued(&entry.0))
.unwrap_or_default()
{
db_entry = self.cursor.next()?;
@ -304,7 +292,7 @@ where
// Compare two entries and return the lowest.
let result = Self::next_slot(post_state_entry, db_entry);
self.last_slot = result.as_ref().map(|entry| entry.key);
self.last_slot = result.as_ref().map(|entry| entry.0);
Ok(result)
}
@ -312,9 +300,9 @@ where
///
/// # Panics
///
/// If the account key is not set. [HashedStorageCursor::seek] must be called first in order to
/// If the account key is not set. [HashedCursor::seek] must be called first in order to
/// position the cursor.
fn next(&mut self) -> Result<Option<StorageEntry>, reth_db::DatabaseError> {
fn next(&mut self) -> Result<Option<(B256, Self::Value)>, reth_db::DatabaseError> {
let last_slot = match self.last_slot.as_ref() {
Some(slot) => slot,
None => return Ok(None), // no previous entry was found
@ -329,7 +317,7 @@ where
// If the entry was already returned or is zero-values, move to the next.
while db_entry
.as_ref()
.map(|entry| &entry.key == last_slot || self.is_slot_zero_valued(&entry.key))
.map(|entry| &entry.0 == last_slot || self.is_slot_zero_valued(&entry.0))
.unwrap_or_default()
{
db_entry = self.cursor.next()?;
@ -350,11 +338,33 @@ where
// Compare two entries and return the lowest.
let result = Self::next_slot(post_state_entry, db_entry);
self.last_slot = result.as_ref().map(|entry| entry.key);
self.last_slot = result.as_ref().map(|entry| entry.0);
Ok(result)
}
}
impl<'b, C> HashedStorageCursor for HashedPostStateStorageCursor<'b, C>
where
C: HashedStorageCursor<Value = U256>,
{
/// Returns `true` if the account has no storage entries.
///
/// This function should be called before attempting to call [HashedCursor::seek] or
/// [HashedCursor::next].
fn is_storage_empty(&mut self) -> Result<bool, reth_db::DatabaseError> {
let is_empty = match self.post_state.storages.get(&self.hashed_address) {
Some(storage) => {
// If the storage has been wiped at any point
storage.wiped &&
// and the current storage does not contain any non-zero values
storage.non_zero_valued_slots.is_empty()
}
None => self.cursor.is_storage_empty()?,
};
Ok(is_empty)
}
}
#[cfg(test)]
mod tests {
use super::*;
@ -363,6 +373,7 @@ mod tests {
use reth_db::{
database::Database, tables, test_utils::create_test_rw_db, transaction::DbTxMut,
};
use reth_primitives::StorageEntry;
use std::collections::BTreeMap;
fn assert_account_cursor_order(
@ -391,11 +402,11 @@ mod tests {
let mut expected_storage = storage.into_iter();
let first_storage = cursor.seek(B256::default()).unwrap();
assert_eq!(first_storage.map(|e| (e.key, e.value)), expected_storage.next());
assert_eq!(first_storage, expected_storage.next());
for expected_entry in expected_storage {
let next_cursor_storage = cursor.next().unwrap();
assert_eq!(next_cursor_storage.map(|e| (e.key, e.value)), Some(expected_entry));
assert_eq!(next_cursor_storage, Some(expected_entry));
}
assert!(cursor.next().unwrap().is_none());

View File

@ -1,10 +1,10 @@
use crate::{
hashed_cursor::{HashedAccountCursor, HashedCursorFactory, HashedStorageCursor},
hashed_cursor::{HashedCursor, HashedCursorFactory, HashedStorageCursor},
trie_cursor::TrieCursor,
walker::TrieWalker,
};
use reth_db::DatabaseError;
use reth_primitives::{trie::Nibbles, Account, StorageEntry, B256, U256};
use reth_primitives::{trie::Nibbles, Account, B256, U256};
/// Represents a branch node in the trie.
#[derive(Debug)]
@ -90,7 +90,7 @@ impl<C, H> AccountNodeIter<C, H> {
impl<C, H> AccountNodeIter<C, H>
where
C: TrieCursor,
H: HashedAccountCursor,
H: HashedCursor<Value = Account>,
{
/// Return the next account trie node to be added to the hash builder.
///
@ -168,7 +168,7 @@ pub struct StorageNodeIter<C, H> {
pub hashed_storage_cursor: H,
/// Current hashed storage entry.
current_hashed_entry: Option<StorageEntry>,
current_hashed_entry: Option<(B256, U256)>,
/// Flag indicating whether we should check the current walker key.
current_walker_key_checked: bool,
}
@ -188,7 +188,7 @@ impl<C, H> StorageNodeIter<C, H> {
impl<C, H> StorageNodeIter<C, H>
where
C: TrieCursor,
H: HashedStorageCursor,
H: HashedStorageCursor<Value = U256>,
{
/// Return the next storage trie node to be added to the hash builder.
///
@ -219,8 +219,7 @@ where
}
// Check for a current hashed storage entry.
if let Some(StorageEntry { key: hashed_key, value }) = self.current_hashed_entry.take()
{
if let Some((hashed_key, value)) = self.current_hashed_entry.take() {
// Compare keys and proceed accordingly.
if self.walker.key().map_or(false, |key| key < &Nibbles::unpack(hashed_key)) {
self.current_walker_key_checked = false;
@ -233,14 +232,15 @@ where
}
// Attempt to get the next unprocessed key from the walker.
if let Some(seek_key) = self.walker.next_unprocessed_key() {
// Seek and update the current hashed entry based on the new seek key.
self.current_hashed_entry = self.hashed_storage_cursor.seek(seek_key)?;
self.walker.advance()?;
} else {
match self.walker.next_unprocessed_key() {
Some(seek_key) => {
// Seek and update the current hashed entry based on the new seek key.
self.current_hashed_entry = self.hashed_storage_cursor.seek(seek_key)?;
self.walker.advance()?;
}
// No more keys to process, break the loop.
break
}
None => break,
};
}
Ok(None) // Return None if no more nodes are available.