chore(trie): return nibbles from TrieCursor::current (#9227)

This commit is contained in:
Roman Krasiuk
2024-07-01 10:39:00 -07:00
committed by GitHub
parent 52068ccee6
commit ad8ec33dc3
7 changed files with 54 additions and 42 deletions

View File

@ -299,7 +299,9 @@ where
};
trie_updates.extend(
walker_deleted_keys.into_iter().map(|key| (key, TrieOp::Delete)),
walker_deleted_keys
.into_iter()
.map(|nibbles| (TrieKey::AccountNode(nibbles), TrieOp::Delete)),
);
trie_updates.extend_with_account_updates(hash_builder_updates);

View File

@ -1,5 +1,5 @@
use super::{TrieCursor, TrieCursorFactory};
use crate::{updates::TrieKey, BranchNodeCompact, Nibbles, StoredNibbles, StoredNibblesSubKey};
use crate::{BranchNodeCompact, Nibbles, StoredNibbles, StoredNibblesSubKey};
use reth_db::{tables, DatabaseError};
use reth_db_api::{
cursor::{DbCursorRO, DbDupCursorRO},
@ -60,8 +60,8 @@ where
}
/// Retrieves the current key in the cursor.
fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
Ok(self.0.current()?.map(|(k, _)| TrieKey::AccountNode(k.0)))
fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
Ok(self.0.current()?.map(|(k, _)| k.0))
}
}
@ -109,8 +109,8 @@ where
}
/// Retrieves the current value in the storage trie cursor.
fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
Ok(self.cursor.current()?.map(|(k, v)| TrieKey::StorageNode(k, v.nibbles.0)))
fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
Ok(self.cursor.current()?.map(|(_, v)| v.nibbles.0))
}
}

View File

@ -42,7 +42,7 @@ impl<'a, CF: TrieCursorFactory> TrieCursorFactory for InMemoryTrieCursorFactory<
pub struct InMemoryAccountTrieCursor<'a, C> {
cursor: C,
trie_updates: &'a TrieUpdatesSorted,
last_key: Option<TrieKey>,
last_key: Option<Nibbles>,
}
impl<'a, C> InMemoryAccountTrieCursor<'a, C> {
@ -56,12 +56,12 @@ impl<'a, C: TrieCursor> TrieCursor for InMemoryAccountTrieCursor<'a, C> {
&mut self,
key: Nibbles,
) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
if let Some((trie_key, trie_op)) = self.trie_updates.find_account_node(&key) {
self.last_key = Some(trie_key);
if let Some((nibbles, trie_op)) = self.trie_updates.find_account_node(&key) {
self.last_key = Some(nibbles);
Ok(trie_op.into_update().map(|node| (key, node)))
} else {
let result = self.cursor.seek_exact(key)?;
self.last_key = result.as_ref().map(|(k, _)| TrieKey::AccountNode(k.clone()));
self.last_key = result.as_ref().map(|(key, _)| key.clone());
Ok(result)
}
}
@ -78,20 +78,22 @@ impl<'a, C: TrieCursor> TrieCursor for InMemoryAccountTrieCursor<'a, C> {
.cloned();
if let Some((trie_key, trie_op)) = trie_update_entry {
let nibbles = match &trie_key {
TrieKey::AccountNode(nibbles) => nibbles.clone(),
let nibbles = match trie_key {
TrieKey::AccountNode(nibbles) => {
self.last_key = Some(nibbles.clone());
nibbles
}
_ => panic!("Invalid trie key"),
};
self.last_key = Some(trie_key);
return Ok(trie_op.into_update().map(|node| (nibbles, node)))
}
let result = self.cursor.seek(key)?;
self.last_key = result.as_ref().map(|(k, _)| TrieKey::AccountNode(k.clone()));
self.last_key = result.as_ref().map(|(key, _)| key.clone());
Ok(result)
}
fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
if self.last_key.is_some() {
Ok(self.last_key.clone())
} else {
@ -108,7 +110,7 @@ pub struct InMemoryStorageTrieCursor<'a, C> {
trie_update_index: usize,
trie_updates: &'a TrieUpdatesSorted,
hashed_address: B256,
last_key: Option<TrieKey>,
last_key: Option<Nibbles>,
}
impl<'a, C> InMemoryStorageTrieCursor<'a, C> {
@ -129,8 +131,7 @@ impl<'a, C: TrieCursor> TrieCursor for InMemoryStorageTrieCursor<'a, C> {
Ok(trie_op.into_update().map(|node| (key, node)))
} else {
let result = self.cursor.seek_exact(key)?;
self.last_key =
result.as_ref().map(|(k, _)| TrieKey::StorageNode(self.hashed_address, k.clone()));
self.last_key = result.as_ref().map(|(key, _)| key.clone());
Ok(result)
}
}
@ -151,20 +152,21 @@ impl<'a, C: TrieCursor> TrieCursor for InMemoryStorageTrieCursor<'a, C> {
trie_update_entry.filter(|(k, _)| matches!(k, TrieKey::StorageNode(_, _)))
{
let nibbles = match trie_key {
TrieKey::StorageNode(_, nibbles) => nibbles.clone(),
TrieKey::StorageNode(_, nibbles) => {
self.last_key = Some(nibbles.clone());
nibbles.clone()
}
_ => panic!("this should not happen!"),
};
self.last_key = Some(trie_key.clone());
return Ok(trie_op.as_update().map(|node| (nibbles, node.clone())))
}
let result = self.cursor.seek(key)?;
self.last_key =
result.as_ref().map(|(k, _)| TrieKey::StorageNode(self.hashed_address, k.clone()));
self.last_key = result.as_ref().map(|(key, _)| key.clone());
Ok(result)
}
fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
if self.last_key.is_some() {
Ok(self.last_key.clone())
} else {

View File

@ -1,4 +1,4 @@
use crate::{updates::TrieKey, BranchNodeCompact, Nibbles};
use crate::{BranchNodeCompact, Nibbles};
use reth_db::DatabaseError;
use reth_primitives::B256;
@ -51,5 +51,5 @@ pub trait TrieCursor: Send + Sync {
-> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError>;
/// Get the current entry.
fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError>;
fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError>;
}

View File

@ -1,5 +1,5 @@
use super::{TrieCursor, TrieCursorFactory};
use crate::{updates::TrieKey, BranchNodeCompact, Nibbles};
use crate::{BranchNodeCompact, Nibbles};
use reth_db::DatabaseError;
use reth_primitives::B256;
@ -49,7 +49,7 @@ impl TrieCursor for NoopAccountTrieCursor {
}
/// Retrieves the current cursor position within the account trie.
fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
Ok(None)
}
}
@ -77,7 +77,7 @@ impl TrieCursor for NoopStorageTrieCursor {
}
/// Retrieves the current cursor position within storage tries.
fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
Ok(None)
}
}

View File

@ -133,7 +133,9 @@ impl TrieUpdates {
) {
// Add updates from trie walker.
let (_, deleted_keys) = walker.split();
self.extend(deleted_keys.into_iter().map(|key| (key, TrieOp::Delete)));
self.extend(
deleted_keys.into_iter().map(|nibbles| (TrieKey::AccountNode(nibbles), TrieOp::Delete)),
);
// Add account node updates from hash builder.
let (_, hash_builder_updates) = hash_builder.split();
@ -154,7 +156,11 @@ impl TrieUpdates {
) {
// Add updates from trie walker.
let (_, deleted_keys) = walker.split();
self.extend(deleted_keys.into_iter().map(|key| (key, TrieOp::Delete)));
self.extend(
deleted_keys
.into_iter()
.map(|nibbles| (TrieKey::StorageNode(hashed_address, nibbles), TrieOp::Delete)),
);
// Add storage node updates from hash builder.
let (_, hash_builder_updates) = hash_builder.split();
@ -248,11 +254,12 @@ pub struct TrieUpdatesSorted {
impl TrieUpdatesSorted {
/// Find the account node with the given nibbles.
pub fn find_account_node(&self, key: &Nibbles) -> Option<(TrieKey, TrieOp)> {
self.trie_operations
.iter()
.find(|(k, _)| matches!(k, TrieKey::AccountNode(nibbles) if nibbles == key))
.cloned()
pub fn find_account_node(&self, key: &Nibbles) -> Option<(Nibbles, TrieOp)> {
self.trie_operations.iter().find_map(|(k, op)| {
k.as_account_node_key()
.filter(|nibbles| nibbles == &key)
.map(|nibbles| (nibbles.clone(), op.clone()))
})
}
/// Find the storage node with the given hashed address and key.
@ -260,9 +267,11 @@ impl TrieUpdatesSorted {
&self,
hashed_address: &B256,
key: &Nibbles,
) -> Option<(TrieKey, TrieOp)> {
self.trie_operations.iter().find(|(k, _)| {
matches!(k, TrieKey::StorageNode(address, nibbles) if address == hashed_address && nibbles == key)
}).cloned()
) -> Option<(Nibbles, TrieOp)> {
self.trie_operations.iter().find_map(|(k, op)| {
k.as_storage_node_key()
.filter(|(address, nibbles)| address == &hashed_address && nibbles == &key)
.map(|(_, nibbles)| (nibbles.clone(), op.clone()))
})
}
}

View File

@ -1,7 +1,6 @@
use crate::{
prefix_set::PrefixSet,
trie_cursor::{CursorSubNode, TrieCursor},
updates::TrieKey,
BranchNodeCompact, Nibbles,
};
use reth_db::DatabaseError;
@ -24,7 +23,7 @@ pub struct TrieWalker<C> {
/// A `PrefixSet` representing the changes to be applied to the trie.
pub changes: PrefixSet,
/// The retained trie node keys that need to be deleted.
deleted_keys: Option<HashSet<TrieKey>>,
deleted_keys: Option<HashSet<Nibbles>>,
}
impl<C> TrieWalker<C> {
@ -45,7 +44,7 @@ impl<C> TrieWalker<C> {
}
/// Split the walker into stack and trie updates.
pub fn split(mut self) -> (Vec<CursorSubNode>, HashSet<TrieKey>) {
pub fn split(mut self) -> (Vec<CursorSubNode>, HashSet<Nibbles>) {
let keys = self.deleted_keys.take();
(self.stack, keys.unwrap_or_default())
}