perf(trie): avoid update reallocation & track wiped (#12929)

crates/trie/sparse/src/state.rs

@@ -22,8 +22,6 @@ pub struct SparseStateTrie {
     storages: HashMap<B256, SparseTrie>,
     /// Collection of revealed account and storage keys.
     revealed: HashMap<B256, HashSet<B256>>,
-    /// Collection of addresses that had their storage tries wiped.
-    wiped_storages: HashSet<B256>,
     /// Flag indicating whether trie updates should be retained.
     retain_updates: bool,
     /// Reusable buffer for RLP encoding of trie accounts.
@@ -36,7 +34,6 @@ impl Default for SparseStateTrie {
             state: Default::default(),
             storages: Default::default(),
             revealed: Default::default(),
-            wiped_storages: Default::default(),
             retain_updates: false,
             account_rlp_buf: Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE),
         }
@@ -275,9 +272,10 @@ impl SparseStateTrie {
 
     /// Wipe the storage trie at the provided address.
     pub fn wipe_storage(&mut self, address: B256) -> SparseStateTrieResult<()> {
-        let Some(trie) = self.storages.get_mut(&address) else { return Ok(()) };
-        self.wiped_storages.insert(address);
-        trie.wipe().map_err(Into::into)
+        if let Some(trie) = self.storages.get_mut(&address) {
+            trie.wipe()?;
+        }
+        Ok(())
     }
 
     /// Calculates the hashes of the nodes below the provided level.
@@ -302,8 +300,8 @@ impl SparseStateTrie {
         self.state.as_revealed_mut().map(|state| {
             let updates = state.take_updates();
             TrieUpdates {
-                account_nodes: HashMap::from_iter(updates.updated_nodes),
-                removed_nodes: HashSet::from_iter(updates.removed_nodes),
+                account_nodes: updates.updated_nodes,
+                removed_nodes: updates.removed_nodes,
                 storage_tries: self
                     .storages
                     .iter_mut()
@@ -311,9 +309,9 @@ impl SparseStateTrie {
                         let trie = trie.as_revealed_mut().unwrap();
                         let updates = trie.take_updates();
                         let updates = StorageTrieUpdates {
-                            is_deleted: self.wiped_storages.contains(address),
-                            storage_nodes: HashMap::from_iter(updates.updated_nodes),
-                            removed_nodes: HashSet::from_iter(updates.removed_nodes),
+                            is_deleted: updates.wiped,
+                            storage_nodes: updates.updated_nodes,
+                            removed_nodes: updates.removed_nodes,
                         };
                         (*address, updates)
                     })

crates/trie/sparse/src/trie.rs

@@ -111,6 +111,7 @@ pub struct RevealedSparseTrie {
     prefix_set: PrefixSetMut,
     /// Reusable buffer for RLP encoding of nodes.
     rlp_buf: Vec<u8>,
+    /// Retained trie updates.
     updates: Option<SparseTrieUpdates>,
 }
 
@@ -607,8 +608,10 @@ impl RevealedSparseTrie {
 
     /// Wipe the trie, removing all values and nodes, and replacing the root with an empty node.
     pub fn wipe(&mut self) {
+        let updates_retained = self.updates.is_some();
         *self = Self::default();
         self.prefix_set = PrefixSetMut::all();
+        self.updates = updates_retained.then(SparseTrieUpdates::wiped);
     }
 
     /// Return the root of the sparse trie.
@@ -1030,12 +1033,18 @@ impl RlpNodeBuffers {
 pub struct SparseTrieUpdates {
     pub(crate) updated_nodes: HashMap<Nibbles, BranchNodeCompact>,
     pub(crate) removed_nodes: HashSet<Nibbles>,
+    pub(crate) wiped: bool,
+}
+
+impl SparseTrieUpdates {
+    /// Create new wiped sparse trie updates.
+    pub fn wiped() -> Self {
+        Self { wiped: true, ..Default::default() }
+    }
 }
 
 #[cfg(test)]
 mod tests {
-    use std::collections::BTreeMap;
-
     use super::*;
     use alloy_primitives::{map::HashSet, U256};
     use alloy_rlp::Encodable;
@@ -1057,6 +1066,7 @@ mod tests {
         proof::{ProofNodes, ProofRetainer},
         HashBuilder,
     };
+    use std::collections::BTreeMap;
 
     /// Pad nibbles to the length of a B256 hash with zeros on the left.
     fn pad_nibbles_left(nibbles: Nibbles) -> Nibbles {