chore: improve HashedStorage (paradigmxyz#6068)

Co-authored-by: Roman Krasiuk <[email protected]>

crates/trie/Cargo.toml

@@ -22,9 +22,6 @@ revm.workspace = true
 alloy-rlp.workspace = true
 alloy-chains.workspace = true
 
-# tokio
-tokio = { workspace = true, default-features = false, features = ["sync"] }
-
 # tracing
 tracing.workspace = true
 

crates/trie/src/hashed_cursor/post_state.rs

@@ -659,7 +659,7 @@ mod tests {
         {
             let wiped = true;
             let mut hashed_storage = HashedStorage::new(wiped);
-            hashed_storage.insert_zero_valued_slot(B256::random());
+            hashed_storage.insert_storage(B256::random(), U256::ZERO);
 
             let mut hashed_post_state = HashedPostState::default();
             hashed_post_state.insert_hashed_storage(address, hashed_storage);
@@ -674,7 +674,7 @@ mod tests {
         {
             let wiped = true;
             let mut hashed_storage = HashedStorage::new(wiped);
-            hashed_storage.insert_non_zero_valued_storage(B256::random(), U256::from(1));
+            hashed_storage.insert_storage(B256::random(), U256::from(1));
 
             let mut hashed_post_state = HashedPostState::default();
             hashed_post_state.insert_hashed_storage(address, hashed_storage);
@@ -710,7 +710,7 @@ mod tests {
         let wiped = false;
         let mut hashed_storage = HashedStorage::new(wiped);
         for (slot, value) in post_state_storage.iter() {
-            hashed_storage.insert_non_zero_valued_storage(*slot, *value);
+            hashed_storage.insert_storage(*slot, *value);
         }
 
         let mut hashed_post_state = HashedPostState::default();
@@ -746,11 +746,7 @@ mod tests {
         let wiped = false;
         let mut hashed_storage = HashedStorage::new(wiped);
         for (slot, value) in post_state_storage.iter() {
-            if value.is_zero() {
-                hashed_storage.insert_zero_valued_slot(*slot);
-            } else {
-                hashed_storage.insert_non_zero_valued_storage(*slot, *value);
-            }
+            hashed_storage.insert_storage(*slot, *value);
         }
 
         let mut hashed_post_state = HashedPostState::default();
@@ -788,7 +784,7 @@ mod tests {
         let wiped = true;
         let mut hashed_storage = HashedStorage::new(wiped);
         for (slot, value) in post_state_storage.iter() {
-            hashed_storage.insert_non_zero_valued_storage(*slot, *value);
+            hashed_storage.insert_storage(*slot, *value);
         }
 
         let mut hashed_post_state = HashedPostState::default();
@@ -823,7 +819,7 @@ mod tests {
         let wiped = false;
         let mut hashed_storage = HashedStorage::new(wiped);
         for (slot, value) in storage.iter() {
-            hashed_storage.insert_non_zero_valued_storage(*slot, *value);
+            hashed_storage.insert_storage(*slot, *value);
         }
 
         let mut hashed_post_state = HashedPostState::default();
@@ -860,11 +856,7 @@ mod tests {
             for (address, (wiped, storage)) in &post_state_storages {
                 let mut hashed_storage = HashedStorage::new(*wiped);
                 for (slot, value) in storage {
-                    if value.is_zero() {
-                        hashed_storage.insert_zero_valued_slot(*slot);
-                    } else {
-                        hashed_storage.insert_non_zero_valued_storage(*slot, *value);
-                    }
+                    hashed_storage.insert_storage(*slot, *value);
                 }
                 hashed_post_state.insert_hashed_storage(*address, hashed_storage);
             }

crates/trie/src/prefix_set/mod.rs

@@ -12,7 +12,7 @@ pub use loader::{LoadedPrefixSets, PrefixSetLoader};
 /// Internally, this implementation uses a `Vec` and aims to act like a `BTreeSet` in being both
 /// sorted and deduplicated. It does this by keeping a `sorted` flag. The `sorted` flag represents
 /// whether or not the `Vec` is definitely sorted. When a new element is added, it is set to
-/// `false.`. The `Vec` is sorted and deduplicated when `sorted` is `false` and:
+/// `false.`. The `Vec` is sorted and deduplicated when `sorted` is `true` and:
 ///  * An element is being checked for inclusion (`contains`), or
 ///  * The set is being converted into an immutable `PrefixSet` (`freeze`)
 ///

crates/trie/src/state.rs

@@ -57,11 +57,7 @@ impl HashedPostState {
 
             for (key, value) in account.storage.iter() {
                 let hashed_key = keccak256(B256::new(key.to_be_bytes()));
-                if value.present_value.is_zero() {
-                    hashed_storage.insert_zero_valued_slot(hashed_key);
-                } else {
-                    hashed_storage.insert_non_zero_valued_storage(hashed_key, value.present_value);
-                }
+                hashed_storage.insert_storage(hashed_key, value.present_value);
             }
             hashed_state.insert_hashed_storage(hashed_address, hashed_storage)
         }
@@ -258,15 +254,14 @@ impl HashedStorage {
         }
     }
 
-    /// Insert non zero-valued storage entry.
-    pub fn insert_non_zero_valued_storage(&mut self, slot: B256, value: U256) {
-        debug_assert!(value != U256::ZERO, "value cannot be zero");
-        self.non_zero_valued_storage.push((slot, value));
-        self.sorted = false;
-    }
-
-    /// Insert zero-valued storage slot.
-    pub fn insert_zero_valued_slot(&mut self, slot: B256) {
-        self.zero_valued_slots.insert(slot);
+    /// Insert storage entry.
+    #[inline]
+    pub fn insert_storage(&mut self, slot: B256, value: U256) {
+        if value.is_zero() {
+            self.zero_valued_slots.insert(slot);
+        } else {
+            self.non_zero_valued_storage.push((slot, value));
+            self.sorted = false;
+        }
     }
 }