diff --git a/crates/trie/parallel/src/async_root.rs b/crates/trie/parallel/src/async_root.rs index a36a01be5e..e1a031a038 100644 --- a/crates/trie/parallel/src/async_root.rs +++ b/crates/trie/parallel/src/async_root.rs @@ -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)?, diff --git a/crates/trie/parallel/src/parallel_root.rs b/crates/trie/parallel/src/parallel_root.rs index edf552096f..e276d7e055 100644 --- a/crates/trie/parallel/src/parallel_root.rs +++ b/crates/trie/parallel/src/parallel_root.rs @@ -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)?, diff --git a/crates/trie/trie/src/trie.rs b/crates/trie/trie/src/trie.rs index 0ae43cdf84..4d23bfc981 100644 --- a/crates/trie/trie/src/trie.rs +++ b/crates/trie/trie/src/trie.rs @@ -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); diff --git a/crates/trie/trie/src/updates.rs b/crates/trie/trie/src/updates.rs index d0027d6585..b982c94559 100644 --- a/crates/trie/trie/src/updates.rs +++ b/crates/trie/trie/src/updates.rs @@ -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) { self.trie_operations.extend(updates); @@ -144,8 +132,8 @@ impl TrieUpdates { destroyed_accounts: HashSet, ) { // 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(); diff --git a/crates/trie/trie/src/walker.rs b/crates/trie/trie/src/walker.rs index 6486a9b080..990b7fb74d 100644 --- a/crates/trie/trie/src/walker.rs +++ b/crates/trie/trie/src/walker.rs @@ -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 { 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, + /// The retained trie node keys that need to be deleted. + deleted_keys: Option>, } impl TrieWalker { /// Constructs a new `TrieWalker` from existing stack and a cursor. pub fn from_stack(cursor: C, stack: Vec, 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, TrieUpdates) { - let trie_updates = self.trie_updates.take(); - (self.stack, trie_updates.unwrap_or_default()) + pub fn split(mut self) -> (Vec, HashSet) { + let keys = self.deleted_keys.take(); + (self.stack, keys.unwrap_or_default()) } /// Prints the current stack of trie nodes. @@ -63,9 +59,9 @@ impl TrieWalker { 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 TrieWalker { 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 TrieWalker { // 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); } }