perf(trie): remove some clones (#10406)

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 }
     }

crates/storage/provider/src/providers/bundle_state_provider.rs

@@ -143,7 +143,7 @@ impl<SP: StateProvider, EDP: ExecutionDataProvider> StateRootProvider
                 )
             })
             .unwrap_or_else(|| HashedStorage::new(false));
-        storage.extend(hashed_storage);
+        storage.extend(&hashed_storage);
         self.state_provider.hashed_storage_root(address, storage)
     }
 }

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<B256> {
         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()))
     }

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<Item = (B256, Cow<'a, HashedStorage>)>,
+    ) {
+        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`].

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<Item = (Nibbles, BranchNodeCompact)>) -> 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<C>(&mut self, walker: TrieWalker<C>, 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>(I);
-
-impl<I: Iterator<Item = Nibbles>> ExcludeEmpty<I> {
-    fn from_iter<T: IntoIterator<Item = Nibbles, IntoIter = I>>(iter: T) -> Self {
-        Self(iter.into_iter())
-    }
-}
-
-impl<I: Iterator<Item = Nibbles>> Iterator for ExcludeEmpty<I> {
-    type Item = Nibbles;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        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<Item = Nibbles>) -> impl Iterator<Item = Nibbles> {
+    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>(I);
-
-impl<V, I: Iterator<Item = (Nibbles, V)>> ExcludeEmptyFromPair<I> {
-    fn from_iter<T: IntoIterator<Item = (Nibbles, V), IntoIter = I>>(iter: T) -> Self {
-        Self(iter.into_iter())
-    }
-}
-
-impl<V, I: Iterator<Item = (Nibbles, V)>> Iterator for ExcludeEmptyFromPair<I> {
-    type Item = (Nibbles, V);
-
-    fn next(&mut self) -> Option<Self::Item> {
-        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<V>(
+    iter: impl IntoIterator<Item = (Nibbles, V)>,
+) -> impl Iterator<Item = (Nibbles, V)> {
+    iter.into_iter().filter(|(n, _)| !n.is_empty())
 }