diff --git a/crates/chain-state/src/memory_overlay.rs b/crates/chain-state/src/memory_overlay.rs index 72d2c8a12ad3..c4a0f288822e 100644 --- a/crates/chain-state/src/memory_overlay.rs +++ b/crates/chain-state/src/memory_overlay.rs @@ -39,8 +39,8 @@ impl MemoryOverlayStateProvider { let mut hashed_post_state = HashedPostState::default(); let mut trie_updates = TrieUpdates::default(); for block in in_memory.iter().rev() { - hashed_post_state.extend(block.hashed_state.as_ref().clone()); - trie_updates.extend(block.trie.as_ref().clone()); + hashed_post_state.extend_ref(block.hashed_state.as_ref()); + trie_updates.extend_ref(block.trie.as_ref()); } Self { in_memory, hashed_post_state, trie_updates, historical } } diff --git a/crates/storage/provider/src/providers/bundle_state_provider.rs b/crates/storage/provider/src/providers/bundle_state_provider.rs index dd4d478c51fb..1359cf4e8cca 100644 --- a/crates/storage/provider/src/providers/bundle_state_provider.rs +++ b/crates/storage/provider/src/providers/bundle_state_provider.rs @@ -143,7 +143,7 @@ impl StateRootProvider ) }) .unwrap_or_else(|| HashedStorage::new(false)); - storage.extend(hashed_storage); + storage.extend(&hashed_storage); self.state_provider.hashed_storage_root(address, storage) } } diff --git a/crates/storage/provider/src/providers/state/historical.rs b/crates/storage/provider/src/providers/state/historical.rs index 2a4c51a718c4..78325721266e 100644 --- a/crates/storage/provider/src/providers/state/historical.rs +++ b/crates/storage/provider/src/providers/state/historical.rs @@ -341,7 +341,7 @@ impl<'b, TX: DbTx> StateRootProvider for HistoricalStateProviderRef<'b, TX> { hashed_storage: HashedStorage, ) -> ProviderResult { let mut revert_storage = self.revert_storage(address)?; - revert_storage.extend(hashed_storage); + revert_storage.extend(&hashed_storage); StorageRoot::overlay_root(self.tx, address, revert_storage) .map_err(|err| ProviderError::Database(err.into())) } diff --git a/crates/trie/trie/src/state.rs b/crates/trie/trie/src/state.rs index b1736e6e94ce..1677e3b4931d 100644 --- a/crates/trie/trie/src/state.rs +++ b/crates/trie/trie/src/state.rs @@ -6,7 +6,10 @@ use itertools::Itertools; use rayon::prelude::{IntoParallelIterator, ParallelIterator}; use reth_primitives::{keccak256, Account, Address, B256, U256}; use revm::db::{states::CacheAccount, AccountStatus, BundleAccount}; -use std::collections::{hash_map, HashMap, HashSet}; +use std::{ + borrow::Cow, + collections::{hash_map, HashMap, HashSet}, +}; /// Representation of in-memory hashed state. #[derive(PartialEq, Eq, Clone, Default, Debug)] @@ -99,17 +102,42 @@ impl HashedPostState { /// Extend this hashed post state with contents of another. /// Entries in the second hashed post state take precedence. pub fn extend(&mut self, other: Self) { - for (hashed_address, account) in other.accounts { - self.accounts.insert(hashed_address, account); + self.extend_inner(Cow::Owned(other)); + } + + /// Extend this hashed post state with contents of another. + /// Entries in the second hashed post state take precedence. + /// + /// Slightly less efficient than [`Self::extend`], but preferred to `extend(other.clone())`. + pub fn extend_ref(&mut self, other: &Self) { + self.extend_inner(Cow::Borrowed(other)); + } + + fn extend_inner(&mut self, other: Cow<'_, Self>) { + self.accounts.extend(other.accounts.iter().map(|(&k, &v)| (k, v))); + + self.storages.reserve(other.storages.len()); + match other { + Cow::Borrowed(other) => { + self.extend_storages(other.storages.iter().map(|(k, v)| (*k, Cow::Borrowed(v)))) + } + Cow::Owned(other) => { + self.extend_storages(other.storages.into_iter().map(|(k, v)| (k, Cow::Owned(v)))) + } } + } - for (hashed_address, storage) in other.storages { + fn extend_storages<'a>( + &mut self, + storages: impl IntoIterator)>, + ) { + for (hashed_address, storage) in storages { match self.storages.entry(hashed_address) { hash_map::Entry::Vacant(entry) => { - entry.insert(storage); + entry.insert(storage.into_owned()); } hash_map::Entry::Occupied(mut entry) => { - entry.get_mut().extend(storage); + entry.get_mut().extend(&storage); } } } @@ -210,14 +238,12 @@ impl HashedStorage { /// Extend hashed storage with contents of other. /// The entries in second hashed storage take precedence. - pub fn extend(&mut self, other: Self) { + pub fn extend(&mut self, other: &Self) { if other.wiped { self.wiped = true; self.storage.clear(); } - for (hashed_slot, value) in other.storage { - self.storage.insert(hashed_slot, value); - } + self.storage.extend(other.storage.iter().map(|(&k, &v)| (k, v))); } /// Converts hashed storage into [`HashedStorageSorted`]. diff --git a/crates/trie/trie/src/updates.rs b/crates/trie/trie/src/updates.rs index 3f9bc73d3d5c..10bf036a95da 100644 --- a/crates/trie/trie/src/updates.rs +++ b/crates/trie/trie/src/updates.rs @@ -36,14 +36,32 @@ impl TrieUpdates { /// Extends the trie updates. pub fn extend(&mut self, other: Self) { - self.account_nodes.retain(|nibbles, _| !other.removed_nodes.contains(nibbles)); - self.account_nodes.extend(ExcludeEmptyFromPair::from_iter(other.account_nodes)); - self.removed_nodes.extend(ExcludeEmpty::from_iter(other.removed_nodes)); + self.extend_common(&other); + self.account_nodes.extend(exclude_empty_from_pair(other.account_nodes)); + self.removed_nodes.extend(exclude_empty(other.removed_nodes)); for (hashed_address, storage_trie) in other.storage_tries { self.storage_tries.entry(hashed_address).or_default().extend(storage_trie); } } + /// Extends the trie updates. + /// + /// Slightly less efficient than [`Self::extend`], but preferred to `extend(other.clone())`. + pub fn extend_ref(&mut self, other: &Self) { + self.extend_common(other); + self.account_nodes.extend(exclude_empty_from_pair( + other.account_nodes.iter().map(|(k, v)| (k.clone(), v.clone())), + )); + self.removed_nodes.extend(exclude_empty(other.removed_nodes.iter().cloned())); + for (hashed_address, storage_trie) in &other.storage_tries { + self.storage_tries.entry(*hashed_address).or_default().extend_ref(storage_trie); + } + } + + fn extend_common(&mut self, other: &Self) { + self.account_nodes.retain(|nibbles, _| !other.removed_nodes.contains(nibbles)); + } + /// Insert storage updates for a given hashed address. pub fn insert_storage_updates( &mut self, @@ -63,11 +81,11 @@ impl TrieUpdates { ) { // Retrieve deleted keys from trie walker. let (_, removed_node_keys) = walker.split(); - self.removed_nodes.extend(ExcludeEmpty::from_iter(removed_node_keys)); + self.removed_nodes.extend(exclude_empty(removed_node_keys)); // Retrieve updated nodes from hash builder. let (_, updated_nodes) = hash_builder.split(); - self.account_nodes.extend(ExcludeEmptyFromPair::from_iter(updated_nodes)); + self.account_nodes.extend(exclude_empty_from_pair(updated_nodes)); // Add deleted storage tries for destroyed accounts. for destroyed in destroyed_accounts { @@ -104,10 +122,7 @@ pub struct StorageTrieUpdates { impl StorageTrieUpdates { /// Creates a new storage trie updates that are not marked as deleted. pub fn new(updates: impl IntoIterator) -> Self { - Self { - storage_nodes: ExcludeEmptyFromPair::from_iter(updates).collect(), - ..Default::default() - } + Self { storage_nodes: exclude_empty_from_pair(updates).collect(), ..Default::default() } } } @@ -153,25 +168,40 @@ impl StorageTrieUpdates { /// Extends storage trie updates. pub fn extend(&mut self, other: Self) { + self.extend_common(&other); + self.storage_nodes.extend(exclude_empty_from_pair(other.storage_nodes)); + self.removed_nodes.extend(exclude_empty(other.removed_nodes)); + } + + /// Extends storage trie updates. + /// + /// Slightly less efficient than [`Self::extend`], but preferred to `extend(other.clone())`. + pub fn extend_ref(&mut self, other: &Self) { + self.extend_common(other); + self.storage_nodes.extend(exclude_empty_from_pair( + other.storage_nodes.iter().map(|(k, v)| (k.clone(), v.clone())), + )); + self.removed_nodes.extend(exclude_empty(other.removed_nodes.iter().cloned())); + } + + fn extend_common(&mut self, other: &Self) { if other.is_deleted { self.storage_nodes.clear(); self.removed_nodes.clear(); } self.is_deleted |= other.is_deleted; self.storage_nodes.retain(|nibbles, _| !other.removed_nodes.contains(nibbles)); - self.storage_nodes.extend(ExcludeEmptyFromPair::from_iter(other.storage_nodes)); - self.removed_nodes.extend(ExcludeEmpty::from_iter(other.removed_nodes)); } /// Finalize storage trie updates for by taking updates from walker and hash builder. pub fn finalize(&mut self, walker: TrieWalker, hash_builder: HashBuilder) { // Retrieve deleted keys from trie walker. let (_, removed_keys) = walker.split(); - self.removed_nodes.extend(ExcludeEmpty::from_iter(removed_keys)); + self.removed_nodes.extend(exclude_empty(removed_keys)); // Retrieve updated nodes from hash builder. let (_, updated_nodes) = hash_builder.split(); - self.storage_nodes.extend(ExcludeEmptyFromPair::from_iter(updated_nodes)); + self.storage_nodes.extend(exclude_empty_from_pair(updated_nodes)); } /// Convert storage trie updates into [`StorageTrieUpdatesSorted`]. @@ -236,46 +266,14 @@ impl StorageTrieUpdatesSorted { } } -// A wrapper iterator to exclude empty nibbles. -struct ExcludeEmpty(I); - -impl> ExcludeEmpty { - fn from_iter>(iter: T) -> Self { - Self(iter.into_iter()) - } -} - -impl> Iterator for ExcludeEmpty { - type Item = Nibbles; - - fn next(&mut self) -> Option { - loop { - let next = self.0.next()?; - if !next.is_empty() { - return Some(next) - } - } - } +/// Excludes empty nibbles from the given iterator. +fn exclude_empty(iter: impl IntoIterator) -> impl Iterator { + iter.into_iter().filter(|n| !n.is_empty()) } -// A wrapper iterator to exclude empty nibbles from pair where nibbles are the key. -struct ExcludeEmptyFromPair(I); - -impl> ExcludeEmptyFromPair { - fn from_iter>(iter: T) -> Self { - Self(iter.into_iter()) - } -} - -impl> Iterator for ExcludeEmptyFromPair { - type Item = (Nibbles, V); - - fn next(&mut self) -> Option { - loop { - let next = self.0.next()?; - if !next.0.is_empty() { - return Some(next) - } - } - } +/// Excludes empty nibbles from the given iterator of pairs where the nibbles are the key. +fn exclude_empty_from_pair( + iter: impl IntoIterator, +) -> impl Iterator { + iter.into_iter().filter(|(n, _)| !n.is_empty()) }