From da03b7989e61f5d3f2b2a576f0c153d996189b05 Mon Sep 17 00:00:00 2001
From: Hai | RISE <150876604+hai-rise@users.noreply.github.com>
Date: Thu, 5 Dec 2024 20:45:43 +0700
Subject: [PATCH] perf(trie): avoid (de)allocating an extra prefix set (#13020)

---
 crates/trie/db/src/storage.rs    |  4 +++-
 crates/trie/parallel/src/root.rs |  3 ++-
 crates/trie/trie/src/trie.rs     | 13 +++++++------
 3 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/crates/trie/db/src/storage.rs b/crates/trie/db/src/storage.rs
index 6a3bbe1b9651..3e40b298fac9 100644
--- a/crates/trie/db/src/storage.rs
+++ b/crates/trie/db/src/storage.rs
@@ -43,6 +43,7 @@ impl<'a, TX: DbTx> DatabaseStorageRoot<'a, TX>
             DatabaseTrieCursorFactory::new(tx),
             DatabaseHashedCursorFactory::new(tx),
             address,
+            Default::default(),
             #[cfg(feature = "metrics")]
             TrieRootMetrics::new(TrieType::Storage),
         )
@@ -53,6 +54,7 @@ impl<'a, TX: DbTx> DatabaseStorageRoot<'a, TX>
             DatabaseTrieCursorFactory::new(tx),
             DatabaseHashedCursorFactory::new(tx),
             hashed_address,
+            Default::default(),
             #[cfg(feature = "metrics")]
             TrieRootMetrics::new(TrieType::Storage),
         )
@@ -70,10 +72,10 @@ impl<'a, TX: DbTx> DatabaseStorageRoot<'a, TX>
             DatabaseTrieCursorFactory::new(tx),
             HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &state_sorted),
             address,
+            prefix_set,
             #[cfg(feature = "metrics")]
             TrieRootMetrics::new(TrieType::Storage),
         )
-        .with_prefix_set(prefix_set)
         .root()
     }
 }
diff --git a/crates/trie/parallel/src/root.rs b/crates/trie/parallel/src/root.rs
index e66d1f782132..2aace02ed9b8 100644
--- a/crates/trie/parallel/src/root.rs
+++ b/crates/trie/parallel/src/root.rs
@@ -121,10 +121,10 @@ where
                         trie_cursor_factory,
                         hashed_state,
                         hashed_address,
+                        prefix_set,
                         #[cfg(feature = "metrics")]
                         metrics,
                     )
-                    .with_prefix_set(prefix_set)
                     .calculate(retain_updates)?)
                 })();
                 let _ = tx.send(result);
@@ -179,6 +179,7 @@ where
                                 trie_cursor_factory.clone(),
                                 hashed_cursor_factory.clone(),
                                 hashed_address,
+                                Default::default(),
                                 #[cfg(feature = "metrics")]
                                 self.metrics.storage_trie.clone(),
                             )
diff --git a/crates/trie/trie/src/trie.rs b/crates/trie/trie/src/trie.rs
index 28517b23e90f..953b8d31fc4f 100644
--- a/crates/trie/trie/src/trie.rs
+++ b/crates/trie/trie/src/trie.rs
@@ -202,15 +202,13 @@ where
                         self.trie_cursor_factory.clone(),
                         self.hashed_cursor_factory.clone(),
                         hashed_address,
-                        #[cfg(feature = "metrics")]
-                        self.metrics.storage_trie.clone(),
-                    )
-                    .with_prefix_set(
                         self.prefix_sets
                             .storage_prefix_sets
                             .get(&hashed_address)
                             .cloned()
                             .unwrap_or_default(),
+                        #[cfg(feature = "metrics")]
+                        self.metrics.storage_trie.clone(),
                     );
 
                     let storage_root = if retain_updates {
@@ -301,29 +299,32 @@ impl<T, H> StorageRoot<T, H> {
         trie_cursor_factory: T,
         hashed_cursor_factory: H,
         address: Address,
+        prefix_set: PrefixSet,
         #[cfg(feature = "metrics")] metrics: TrieRootMetrics,
     ) -> Self {
         Self::new_hashed(
             trie_cursor_factory,
             hashed_cursor_factory,
             keccak256(address),
+            prefix_set,
             #[cfg(feature = "metrics")]
             metrics,
         )
     }
 
     /// Creates a new storage root calculator given a hashed address.
-    pub fn new_hashed(
+    pub const fn new_hashed(
         trie_cursor_factory: T,
         hashed_cursor_factory: H,
         hashed_address: B256,
+        prefix_set: PrefixSet,
         #[cfg(feature = "metrics")] metrics: TrieRootMetrics,
     ) -> Self {
         Self {
             trie_cursor_factory,
             hashed_cursor_factory,
             hashed_address,
-            prefix_set: PrefixSet::default(),
+            prefix_set,
             #[cfg(feature = "metrics")]
             metrics,
         }