chore(trie): store only deleted keys in TrieWalker (#9226)

This commit is contained in:
Roman Krasiuk
2024-07-01 09:06:16 -07:00
committed by GitHub
parent 20713e9785
commit 898d17bb91
5 changed files with 33 additions and 46 deletions

View File

@@ -132,7 +132,7 @@ where
trie_cursor_factory.account_trie_cursor().map_err(ProviderError::Database)?,
prefix_sets.account_prefix_set,
)
.with_updates(retain_updates);
.with_deletions_retained(retain_updates);
let mut account_node_iter = TrieNodeIter::new(
walker,
hashed_cursor_factory.hashed_account_cursor().map_err(ProviderError::Database)?,

View File

@@ -116,7 +116,7 @@ where
trie_cursor_factory.account_trie_cursor().map_err(ProviderError::Database)?,
prefix_sets.account_prefix_set,
)
.with_updates(retain_updates);
.with_deletions_retained(retain_updates);
let mut account_node_iter = TrieNodeIter::new(
walker,
hashed_cursor_factory.hashed_account_cursor().map_err(ProviderError::Database)?,

View File

@@ -221,7 +221,7 @@ where
state.walker_stack,
self.prefix_sets.account_prefix_set,
)
.with_updates(retain_updates);
.with_deletions_retained(retain_updates);
let node_iter = TrieNodeIter::new(walker, hashed_account_cursor)
.with_last_hashed_key(state.last_account_key);
(hash_builder, node_iter)
@@ -229,7 +229,7 @@ where
None => {
let hash_builder = HashBuilder::default().with_updates(retain_updates);
let walker = TrieWalker::new(trie_cursor, self.prefix_sets.account_prefix_set)
.with_updates(retain_updates);
.with_deletions_retained(retain_updates);
let node_iter = TrieNodeIter::new(walker, hashed_account_cursor);
(hash_builder, node_iter)
}
@@ -286,10 +286,10 @@ where
// Decide if we need to return intermediate progress.
let total_updates_len = trie_updates.len() +
account_node_iter.walker.updates_len() +
account_node_iter.walker.deleted_keys_len() +
hash_builder.updates_len();
if retain_updates && total_updates_len as u64 >= self.threshold {
let (walker_stack, walker_updates) = account_node_iter.walker.split();
let (walker_stack, walker_deleted_keys) = account_node_iter.walker.split();
let (hash_builder, hash_builder_updates) = hash_builder.split();
let state = IntermediateStateRootState {
@@ -298,7 +298,9 @@ where
last_account_key: hashed_address,
};
trie_updates.extend(walker_updates);
trie_updates.extend(
walker_deleted_keys.into_iter().map(|key| (key, TrieOp::Delete)),
);
trie_updates.extend_with_account_updates(hash_builder_updates);
return Ok(StateRootProgress::Progress(
@@ -492,7 +494,8 @@ where
let mut tracker = TrieTracker::default();
let trie_cursor = self.trie_cursor_factory.storage_trie_cursor(self.hashed_address)?;
let walker = TrieWalker::new(trie_cursor, self.prefix_set).with_updates(retain_updates);
let walker =
TrieWalker::new(trie_cursor, self.prefix_set).with_deletions_retained(retain_updates);
let mut hash_builder = HashBuilder::default().with_updates(retain_updates);

View File

@@ -110,18 +110,6 @@ impl IntoIterator for TrieUpdates {
}
impl TrieUpdates {
/// Schedule a delete operation on a trie key.
///
/// # Panics
///
/// If the key already exists and the operation is an update.
pub fn schedule_delete(&mut self, key: TrieKey) {
let existing = self.trie_operations.insert(key, TrieOp::Delete);
if let Some(op) = existing {
assert!(!op.is_update(), "Tried to delete a node that was already updated");
}
}
/// Extend the updates with trie updates.
pub fn extend(&mut self, updates: impl IntoIterator<Item = (TrieKey, TrieOp)>) {
self.trie_operations.extend(updates);
@@ -144,8 +132,8 @@ impl TrieUpdates {
destroyed_accounts: HashSet<B256>,
) {
// Add updates from trie walker.
let (_, walker_updates) = walker.split();
self.extend(walker_updates);
let (_, deleted_keys) = walker.split();
self.extend(deleted_keys.into_iter().map(|key| (key, TrieOp::Delete)));
// Add account node updates from hash builder.
let (_, hash_builder_updates) = hash_builder.split();
@@ -165,8 +153,8 @@ impl TrieUpdates {
hash_builder: HashBuilder,
) {
// Add updates from trie walker.
let (_, walker_updates) = walker.split();
self.extend(walker_updates);
let (_, deleted_keys) = walker.split();
self.extend(deleted_keys.into_iter().map(|key| (key, TrieOp::Delete)));
// Add storage node updates from hash builder.
let (_, hash_builder_updates) = hash_builder.split();

View File

@@ -1,11 +1,12 @@
use crate::{
prefix_set::PrefixSet,
trie_cursor::{CursorSubNode, TrieCursor},
updates::TrieUpdates,
updates::TrieKey,
BranchNodeCompact, Nibbles,
};
use reth_db::DatabaseError;
use reth_primitives::B256;
use std::collections::HashSet;
/// `TrieWalker` is a structure that enables traversal of a Merkle trie.
/// It allows moving through the trie in a depth-first manner, skipping certain branches
@@ -22,36 +23,31 @@ pub struct TrieWalker<C> {
pub can_skip_current_node: bool,
/// A `PrefixSet` representing the changes to be applied to the trie.
pub changes: PrefixSet,
/// The trie updates to be applied to the trie.
trie_updates: Option<TrieUpdates>,
/// The retained trie node keys that need to be deleted.
deleted_keys: Option<HashSet<TrieKey>>,
}
impl<C> TrieWalker<C> {
/// Constructs a new `TrieWalker` from existing stack and a cursor.
pub fn from_stack(cursor: C, stack: Vec<CursorSubNode>, changes: PrefixSet) -> Self {
let mut this =
Self { cursor, changes, stack, can_skip_current_node: false, trie_updates: None };
Self { cursor, changes, stack, can_skip_current_node: false, deleted_keys: None };
this.update_skip_node();
this
}
/// Sets the flag whether the trie updates should be stored.
pub fn with_updates(mut self, retain_updates: bool) -> Self {
self.set_updates(retain_updates);
pub fn with_deletions_retained(mut self, retained: bool) -> Self {
if retained {
self.deleted_keys = Some(HashSet::default());
}
self
}
/// Sets the flag whether the trie updates should be stored.
pub fn set_updates(&mut self, retain_updates: bool) {
if retain_updates {
self.trie_updates = Some(TrieUpdates::default());
}
}
/// Split the walker into stack and trie updates.
pub fn split(mut self) -> (Vec<CursorSubNode>, TrieUpdates) {
let trie_updates = self.trie_updates.take();
(self.stack, trie_updates.unwrap_or_default())
pub fn split(mut self) -> (Vec<CursorSubNode>, HashSet<TrieKey>) {
let keys = self.deleted_keys.take();
(self.stack, keys.unwrap_or_default())
}
/// Prints the current stack of trie nodes.
@@ -63,9 +59,9 @@ impl<C> TrieWalker<C> {
println!("====================== END STACK ======================\n");
}
/// The current length of the trie updates.
pub fn updates_len(&self) -> usize {
self.trie_updates.as_ref().map(|u| u.len()).unwrap_or(0)
/// The current length of the deleted keys.
pub fn deleted_keys_len(&self) -> usize {
self.deleted_keys.as_ref().map_or(0, |u| u.len())
}
/// Returns the current key in the trie.
@@ -117,7 +113,7 @@ impl<C: TrieCursor> TrieWalker<C> {
changes,
stack: vec![CursorSubNode::default()],
can_skip_current_node: false,
trie_updates: None,
deleted_keys: None,
};
// Set up the root node of the trie in the stack, if it exists.
@@ -193,8 +189,8 @@ impl<C: TrieCursor> TrieWalker<C> {
// Delete the current node if it's included in the prefix set or it doesn't contain the root
// hash.
if !self.can_skip_current_node || nibble != -1 {
if let Some((updates, key)) = self.trie_updates.as_mut().zip(self.cursor.current()?) {
updates.schedule_delete(key);
if let Some((keys, key)) = self.deleted_keys.as_mut().zip(self.cursor.current()?) {
keys.insert(key);
}
}