fix(trie): consider all intermediate nodes invalidated for wiped storage

crates/storage/provider/src/writer/mod.rs

@@ -511,7 +511,9 @@ where
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::{test_utils::create_test_provider_factory, AccountReader, TrieWriter};
+    use crate::{
+        test_utils::create_test_provider_factory, AccountReader, StorageTrieWriter, TrieWriter,
+    };
     use reth_db::tables;
     use reth_db_api::{
         cursor::{DbCursorRO, DbCursorRW, DbDupCursorRO},
@@ -521,8 +523,11 @@ mod tests {
     use reth_primitives::{
         keccak256, Account, Address, Receipt, Receipts, StorageEntry, B256, U256,
     };
-    use reth_trie::{test_utils::state_root, HashedPostState, HashedStorage, StateRoot};
-    use reth_trie_db::DatabaseStateRoot;
+    use reth_trie::{
+        test_utils::{state_root, storage_root_prehashed},
+        HashedPostState, HashedStorage, StateRoot, StorageRoot,
+    };
+    use reth_trie_db::{DatabaseStateRoot, DatabaseStorageRoot};
     use revm::{
         db::{
             states::{
@@ -535,7 +540,10 @@ mod tests {
         },
         DatabaseCommit, State,
     };
-    use std::collections::{BTreeMap, HashMap};
+    use std::{
+        collections::{BTreeMap, HashMap},
+        str::FromStr,
+    };
 
     #[test]
     fn wiped_entries_are_removed() {
@@ -1576,4 +1584,55 @@ mod tests {
         // account2 got inserted
         assert_eq!(end_state.state.get(&address2).unwrap().info, Some(account2));
     }
+
+    #[test]
+    fn hashed_state_storage_root() {
+        let address = Address::random();
+        let hashed_address = keccak256(address);
+        let provider_factory = create_test_provider_factory();
+        let provider_rw = provider_factory.provider_rw().unwrap();
+        let tx = provider_rw.tx_ref();
+
+        // insert initial account storage
+        let init_storage = HashedStorage::from_iter(
+            false,
+            [
+                "50000000000000000000000000000004253371b55351a08cb3267d4d265530b6",
+                "512428ed685fff57294d1a9cbb147b18ae5db9cf6ae4b312fa1946ba0561882e",
+                "51e6784c736ef8548f856909870b38e49ef7a4e3e77e5e945e0d5e6fcaa3037f",
+            ]
+            .into_iter()
+            .map(|str| (B256::from_str(str).unwrap(), U256::from(1))),
+        );
+        let mut state = HashedPostState::default();
+        state.storages.insert(hashed_address, init_storage.clone());
+        provider_rw.write_hashed_state(&state.clone().into_sorted()).unwrap();
+
+        // calculate database storage root and write intermediate storage nodes.
+        let (storage_root, _, storage_updates) =
+            StorageRoot::from_tx_hashed(tx, hashed_address).calculate(true).unwrap();
+        assert_eq!(storage_root, storage_root_prehashed(init_storage.storage));
+        assert!(!storage_updates.is_empty());
+        provider_rw
+            .write_individual_storage_trie_updates(hashed_address, &storage_updates)
+            .unwrap();
+
+        // destroy the storage and re-create with new slots
+        let updated_storage = HashedStorage::from_iter(
+            true,
+            [
+                "00deb8486ad8edccfdedfc07109b3667b38a03a8009271aac250cce062d90917",
+                "88d233b7380bb1bcdc866f6871c94685848f54cf0ee033b1480310b4ddb75fc9",
+            ]
+            .into_iter()
+            .map(|str| (B256::from_str(str).unwrap(), U256::from(1))),
+        );
+        let mut state = HashedPostState::default();
+        state.storages.insert(hashed_address, updated_storage.clone());
+        provider_rw.write_hashed_state(&state.clone().into_sorted()).unwrap();
+
+        // re-calculate database storage root
+        let storage_root = StorageRoot::overlay_root(tx, address, updated_storage.clone()).unwrap();
+        assert_eq!(storage_root, storage_root_prehashed(updated_storage.storage));
+    }
 }

crates/trie/trie/src/prefix_set.rs

@@ -84,6 +84,9 @@ pub struct TriePrefixSets {
 /// ```
 #[derive(Clone, Default, Debug)]
 pub struct PrefixSetMut {
+    /// Flag indicating that any entry should be considered changed.
+    /// If set, the keys will be discarded.
+    all: bool,
     keys: Vec<Nibbles>,
 }
 
@@ -92,14 +95,19 @@ where
     I: IntoIterator<Item = Nibbles>,
 {
     fn from(value: I) -> Self {
-        Self { keys: value.into_iter().collect() }
+        Self { all: false, keys: value.into_iter().collect() }
     }
 }
 
 impl PrefixSetMut {
     /// Create [`PrefixSetMut`] with pre-allocated capacity.
     pub fn with_capacity(capacity: usize) -> Self {
-        Self { keys: Vec::with_capacity(capacity) }
+        Self { all: false, keys: Vec::with_capacity(capacity) }
+    }
+
+    /// Create [`PrefixSetMut`] that considers all key changed.
+    pub const fn all() -> Self {
+        Self { all: true, keys: Vec::new() }
     }
 
     /// Inserts the given `nibbles` into the set.
@@ -129,12 +137,16 @@ impl PrefixSetMut {
     ///
     /// If not yet sorted, the elements will be sorted and deduplicated.
     pub fn freeze(mut self) -> PrefixSet {
-        self.keys.sort();
-        self.keys.dedup();
-        // we need to shrink in both the sorted and non-sorted cases because deduping may have
-        // occurred either on `freeze`, or during `contains`.
-        self.keys.shrink_to_fit();
-        PrefixSet { keys: Arc::new(self.keys), index: 0 }
+        if self.all {
+            PrefixSet { index: 0, all: true, keys: Arc::new(Vec::new()) }
+        } else {
+            self.keys.sort();
+            self.keys.dedup();
+            // we need to shrink in both the sorted and non-sorted cases because deduping may have
+            // occurred either on `freeze`, or during `contains`.
+            self.keys.shrink_to_fit();
+            PrefixSet { index: 0, all: false, keys: Arc::new(self.keys) }
+        }
     }
 }
 
@@ -143,15 +155,21 @@ impl PrefixSetMut {
 /// See also [`PrefixSetMut::freeze`].
 #[derive(Debug, Default, Clone)]
 pub struct PrefixSet {
-    keys: Arc<Vec<Nibbles>>,
+    /// Flag indicating that any entry should be considered changed.
+    all: bool,
     index: usize,
+    keys: Arc<Vec<Nibbles>>,
 }
 
 impl PrefixSet {
     /// Returns `true` if any of the keys in the set has the given prefix or
     /// if the given prefix is a prefix of any key in the set.
     #[inline]
     pub fn contains(&mut self, prefix: &[u8]) -> bool {
+        if self.all {
+            return true
+        }
+
         while self.index > 0 && &self.keys[self.index] > prefix {
             self.index -= 1;
         }

crates/trie/trie/src/state.rs

@@ -197,11 +197,15 @@ impl HashedStorage {
 
     /// Construct [`PrefixSetMut`] from hashed storage.
     pub fn construct_prefix_set(&self) -> PrefixSetMut {
-        let mut prefix_set = PrefixSetMut::with_capacity(self.storage.len());
-        for hashed_slot in self.storage.keys() {
-            prefix_set.insert(Nibbles::unpack(hashed_slot));
+        if self.wiped {
+            PrefixSetMut::all()
+        } else {
+            let mut prefix_set = PrefixSetMut::with_capacity(self.storage.len());
+            for hashed_slot in self.storage.keys() {
+                prefix_set.insert(Nibbles::unpack(hashed_slot));
+            }
+            prefix_set
         }
-        prefix_set
     }
 
     /// Extend hashed storage with contents of other.