From c772fe201dd9a4c060d8ce70886b983993d1247d Mon Sep 17 00:00:00 2001
From: Roman Krasiuk <rokrassyuk@gmail.com>
Date: Wed, 27 Nov 2024 23:34:38 +0100
Subject: [PATCH] perf(trie): use std for sparse trie updates & track wiped

---
 crates/trie/sparse/src/state.rs | 22 +++++++++++-----------
 crates/trie/sparse/src/trie.rs  | 24 +++++++++++++++---------
 2 files changed, 26 insertions(+), 20 deletions(-)

diff --git a/crates/trie/sparse/src/state.rs b/crates/trie/sparse/src/state.rs
index 551a47ce2bbfb..25b81d6cff6e1 100644
--- a/crates/trie/sparse/src/state.rs
+++ b/crates/trie/sparse/src/state.rs
@@ -13,15 +13,14 @@ use std::iter::Peekable;
 /// Sparse state trie representing lazy-loaded Ethereum state trie.
 #[derive(Default, Debug)]
 pub struct SparseStateTrie {
-    retain_updates: bool,
     /// Sparse account trie.
     state: SparseTrie,
     /// Sparse storage tries.
     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 update should be retained.
+    retain_updates: bool,
 }
 
 impl SparseStateTrie {
@@ -225,9 +224,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.
@@ -252,8 +252,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()
@@ -261,9 +261,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)
                     })
diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs
index 21f1cf410aa67..e97e5a88f4fa9 100644
--- a/crates/trie/sparse/src/trie.rs
+++ b/crates/trie/sparse/src/trie.rs
@@ -1,9 +1,5 @@
 use crate::{SparseTrieError, SparseTrieResult};
-use alloy_primitives::{
-    hex, keccak256,
-    map::{HashMap, HashSet},
-    B256,
-};
+use alloy_primitives::{hex, keccak256, map::HashMap, B256};
 use alloy_rlp::Decodable;
 use reth_tracing::tracing::debug;
 use reth_trie_common::{
@@ -111,6 +107,7 @@ pub struct RevealedSparseTrie {
     prefix_set: PrefixSetMut,
     /// Reusable buffer for RLP encoding of nodes.
     rlp_buf: Vec<u8>,
+    /// Retained trie updates.
     updates: Option<SparseTrieUpdates>,
 }
 
@@ -602,8 +599,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.
@@ -1023,14 +1022,20 @@ impl RlpNodeBuffers {
 /// The aggregation of sparse trie updates.
 #[derive(Debug, Clone, Default, PartialEq, Eq)]
 pub struct SparseTrieUpdates {
-    pub(crate) updated_nodes: HashMap<Nibbles, BranchNodeCompact>,
-    pub(crate) removed_nodes: HashSet<Nibbles>,
+    pub(crate) updated_nodes: std::collections::HashMap<Nibbles, BranchNodeCompact>,
+    pub(crate) removed_nodes: std::collections::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;
@@ -1052,6 +1057,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 {