From 5706e03422f1d409491297e5a20946d0fd6413d8 Mon Sep 17 00:00:00 2001
From: Roman Krasiuk <rokrassyuk@gmail.com>
Date: Fri, 27 Sep 2024 11:14:09 +0200
Subject: [PATCH] chore(trie): early return on empty state (#11271)

---
 crates/trie/trie/src/state.rs   | 55 ++++++++++++++++++---------------
 crates/trie/trie/src/witness.rs |  5 +++
 2 files changed, 35 insertions(+), 25 deletions(-)

diff --git a/crates/trie/trie/src/state.rs b/crates/trie/trie/src/state.rs
index d634f05f0f..3b0af5cd87 100644
--- a/crates/trie/trie/src/state.rs
+++ b/crates/trie/trie/src/state.rs
@@ -100,6 +100,36 @@ impl HashedPostState {
         self
     }
 
+    /// Returns `true` if the hashed state is empty.
+    pub fn is_empty(&self) -> bool {
+        self.accounts.is_empty() && self.storages.is_empty()
+    }
+
+    /// Construct [`TriePrefixSetsMut`] from hashed post state.
+    /// The prefix sets contain the hashed account and storage keys that have been changed in the
+    /// post state.
+    pub fn construct_prefix_sets(&self) -> TriePrefixSetsMut {
+        // Populate account prefix set.
+        let mut account_prefix_set = PrefixSetMut::with_capacity(self.accounts.len());
+        let mut destroyed_accounts = HashSet::default();
+        for (hashed_address, account) in &self.accounts {
+            account_prefix_set.insert(Nibbles::unpack(hashed_address));
+
+            if account.is_none() {
+                destroyed_accounts.insert(*hashed_address);
+            }
+        }
+
+        // Populate storage prefix sets.
+        let mut storage_prefix_sets = HashMap::with_capacity(self.storages.len());
+        for (hashed_address, hashed_storage) in &self.storages {
+            account_prefix_set.insert(Nibbles::unpack(hashed_address));
+            storage_prefix_sets.insert(*hashed_address, hashed_storage.construct_prefix_set());
+        }
+
+        TriePrefixSetsMut { account_prefix_set, storage_prefix_sets, destroyed_accounts }
+    }
+
     /// 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) {
@@ -166,31 +196,6 @@ impl HashedPostState {
 
         HashedPostStateSorted { accounts, storages }
     }
-
-    /// Construct [`TriePrefixSetsMut`] from hashed post state.
-    /// The prefix sets contain the hashed account and storage keys that have been changed in the
-    /// post state.
-    pub fn construct_prefix_sets(&self) -> TriePrefixSetsMut {
-        // Populate account prefix set.
-        let mut account_prefix_set = PrefixSetMut::with_capacity(self.accounts.len());
-        let mut destroyed_accounts = HashSet::default();
-        for (hashed_address, account) in &self.accounts {
-            account_prefix_set.insert(Nibbles::unpack(hashed_address));
-
-            if account.is_none() {
-                destroyed_accounts.insert(*hashed_address);
-            }
-        }
-
-        // Populate storage prefix sets.
-        let mut storage_prefix_sets = HashMap::with_capacity(self.storages.len());
-        for (hashed_address, hashed_storage) in &self.storages {
-            account_prefix_set.insert(Nibbles::unpack(hashed_address));
-            storage_prefix_sets.insert(*hashed_address, hashed_storage.construct_prefix_set());
-        }
-
-        TriePrefixSetsMut { account_prefix_set, storage_prefix_sets, destroyed_accounts }
-    }
 }
 
 /// Representation of in-memory hashed storage.
diff --git a/crates/trie/trie/src/witness.rs b/crates/trie/trie/src/witness.rs
index 61576aabe3..ef5b358d31 100644
--- a/crates/trie/trie/src/witness.rs
+++ b/crates/trie/trie/src/witness.rs
@@ -83,6 +83,10 @@ where
         mut self,
         state: HashedPostState,
     ) -> Result<HashMap<B256, Bytes>, TrieWitnessError> {
+        if state.is_empty() {
+            return Ok(self.witness)
+        }
+
         let proof_targets = HashMap::from_iter(
             state
                 .accounts
@@ -92,6 +96,7 @@ where
                     (*hashed_address, storage.storage.keys().copied().collect())
                 })),
         );
+
         let mut account_multiproof =
             Proof::new(self.trie_cursor_factory.clone(), self.hashed_cursor_factory.clone())
                 .with_prefix_sets_mut(self.prefix_sets.clone())