diff --git a/crates/storage/provider/src/writer/mod.rs b/crates/storage/provider/src/writer/mod.rs index 24a4a4fc24a5..3d05aa09a44d 100644 --- a/crates/storage/provider/src/writer/mod.rs +++ b/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)); + } } diff --git a/crates/trie/trie/src/prefix_set.rs b/crates/trie/trie/src/prefix_set.rs index 325697c259ed..fa14464241ea 100644 --- a/crates/trie/trie/src/prefix_set.rs +++ b/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, } @@ -92,14 +95,19 @@ where I: IntoIterator, { 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,8 +155,10 @@ impl PrefixSetMut { /// See also [`PrefixSetMut::freeze`]. #[derive(Debug, Default, Clone)] pub struct PrefixSet { - keys: Arc>, + /// Flag indicating that any entry should be considered changed. + all: bool, index: usize, + keys: Arc>, } impl PrefixSet { @@ -152,6 +166,10 @@ impl PrefixSet { /// 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; } diff --git a/crates/trie/trie/src/state.rs b/crates/trie/trie/src/state.rs index 679cefc15add..b1736e6e94ce 100644 --- a/crates/trie/trie/src/state.rs +++ b/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.