From 2cd5f867256ebf63ba5444356ac84748fc1ab605 Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Thu, 4 Jul 2024 15:49:45 +0200 Subject: [PATCH 1/2] chore(trie): return mutable prefix sets from `HashedPostState::construct_prefix_sets` --- crates/blockchain-tree/src/blockchain_tree.rs | 2 +- crates/trie/parallel/benches/root.rs | 2 +- crates/trie/parallel/src/async_root.rs | 2 +- crates/trie/parallel/src/parallel_root.rs | 2 +- crates/trie/trie/src/prefix_set/mod.rs | 29 +++++++++++++++++++ crates/trie/trie/src/state.rs | 18 +++++------- 6 files changed, 40 insertions(+), 15 deletions(-) diff --git a/crates/blockchain-tree/src/blockchain_tree.rs b/crates/blockchain-tree/src/blockchain_tree.rs index da0a147acceb..5974d1a030a0 100644 --- a/crates/blockchain-tree/src/blockchain_tree.rs +++ b/crates/blockchain-tree/src/blockchain_tree.rs @@ -1866,7 +1866,7 @@ mod tests { ); let provider = tree.externals.provider_factory.provider().unwrap(); - let prefix_sets = exec5.hash_state_slow().construct_prefix_sets(); + let prefix_sets = exec5.hash_state_slow().construct_prefix_sets().freeze(); let state_root = StateRoot::from_tx(provider.tx_ref()).with_prefix_sets(prefix_sets).root().unwrap(); assert_eq!(state_root, block5.state_root); diff --git a/crates/trie/parallel/benches/root.rs b/crates/trie/parallel/benches/root.rs index 288e930bae3c..6a7d7a81cc37 100644 --- a/crates/trie/parallel/benches/root.rs +++ b/crates/trie/parallel/benches/root.rs @@ -41,7 +41,7 @@ pub fn calculate_state_root(c: &mut Criterion) { b.to_async(&runtime).iter_with_setup( || { let sorted_state = updated_state.clone().into_sorted(); - let prefix_sets = updated_state.construct_prefix_sets(); + let prefix_sets = updated_state.construct_prefix_sets().freeze(); let provider = provider_factory.provider().unwrap(); (provider, sorted_state, prefix_sets) }, diff --git a/crates/trie/parallel/src/async_root.rs b/crates/trie/parallel/src/async_root.rs index e568be81b4c5..db6152b6a2cf 100644 --- a/crates/trie/parallel/src/async_root.rs +++ b/crates/trie/parallel/src/async_root.rs @@ -86,7 +86,7 @@ where retain_updates: bool, ) -> Result<(B256, TrieUpdates), AsyncStateRootError> { let mut tracker = ParallelTrieTracker::default(); - let prefix_sets = self.hashed_state.construct_prefix_sets(); + let prefix_sets = self.hashed_state.construct_prefix_sets().freeze(); let storage_root_targets = StorageRootTargets::new( self.hashed_state.accounts.keys().copied(), prefix_sets.storage_prefix_sets, diff --git a/crates/trie/parallel/src/parallel_root.rs b/crates/trie/parallel/src/parallel_root.rs index 5e26b97b672b..0983fd47e5a3 100644 --- a/crates/trie/parallel/src/parallel_root.rs +++ b/crates/trie/parallel/src/parallel_root.rs @@ -77,7 +77,7 @@ where retain_updates: bool, ) -> Result<(B256, TrieUpdates), ParallelStateRootError> { let mut tracker = ParallelTrieTracker::default(); - let prefix_sets = self.hashed_state.construct_prefix_sets(); + let prefix_sets = self.hashed_state.construct_prefix_sets().freeze(); let storage_root_targets = StorageRootTargets::new( self.hashed_state.accounts.keys().copied(), prefix_sets.storage_prefix_sets, diff --git a/crates/trie/trie/src/prefix_set/mod.rs b/crates/trie/trie/src/prefix_set/mod.rs index 228e0abee3c0..9d08ad36a98d 100644 --- a/crates/trie/trie/src/prefix_set/mod.rs +++ b/crates/trie/trie/src/prefix_set/mod.rs @@ -8,6 +8,35 @@ use std::{ mod loader; pub use loader::PrefixSetLoader; +/// Collection of mutable prefix sets. +#[derive(Default, Debug)] +pub struct TriePrefixSetsMut { + /// A set of account prefixes that have changed. + pub account_prefix_set: PrefixSetMut, + /// A map containing storage changes with the hashed address as key and a set of storage key + /// prefixes as the value. + pub storage_prefix_sets: HashMap, + /// A set of hashed addresses of destroyed accounts. + pub destroyed_accounts: HashSet, +} + +impl TriePrefixSetsMut { + /// Returns a `TriePrefixSets` with the same elements as these sets. + /// + /// If not yet sorted, the elements will be sorted and deduplicated. + pub fn freeze(self) -> TriePrefixSets { + TriePrefixSets { + account_prefix_set: self.account_prefix_set.freeze(), + storage_prefix_sets: self + .storage_prefix_sets + .into_iter() + .map(|(hashed_address, prefix_set)| (hashed_address, prefix_set.freeze())) + .collect(), + destroyed_accounts: self.destroyed_accounts, + } + } +} + /// Collection of trie prefix sets. #[derive(Default, Debug)] pub struct TriePrefixSets { diff --git a/crates/trie/trie/src/state.rs b/crates/trie/trie/src/state.rs index de7ecc236dbe..cff265f09a99 100644 --- a/crates/trie/trie/src/state.rs +++ b/crates/trie/trie/src/state.rs @@ -1,6 +1,6 @@ use crate::{ hashed_cursor::HashedPostStateCursorFactory, - prefix_set::{PrefixSetMut, TriePrefixSets}, + prefix_set::{PrefixSetMut, TriePrefixSets, TriePrefixSetsMut}, updates::TrieUpdates, Nibbles, StateRoot, }; @@ -170,10 +170,10 @@ impl HashedPostState { HashedPostStateSorted { accounts, storages } } - /// Construct [`TriePrefixSets`] from hashed post state. + /// 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) -> TriePrefixSets { + 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(); @@ -194,14 +194,10 @@ impl HashedPostState { for hashed_slot in hashed_storage.storage.keys() { prefix_set.insert(Nibbles::unpack(hashed_slot)); } - storage_prefix_sets.insert(*hashed_address, prefix_set.freeze()); + storage_prefix_sets.insert(*hashed_address, prefix_set); } - TriePrefixSets { - account_prefix_set: account_prefix_set.freeze(), - storage_prefix_sets, - destroyed_accounts, - } + TriePrefixSetsMut { account_prefix_set, storage_prefix_sets, destroyed_accounts } } /// Calculate the state root for this [`HashedPostState`]. @@ -236,7 +232,7 @@ impl HashedPostState { /// The state root for this [`HashedPostState`]. pub fn state_root(&self, tx: &TX) -> Result { let sorted = self.clone().into_sorted(); - let prefix_sets = self.construct_prefix_sets(); + let prefix_sets = self.construct_prefix_sets().freeze(); StateRoot::from_tx(tx) .with_hashed_cursor_factory(HashedPostStateCursorFactory::new(tx, &sorted)) .with_prefix_sets(prefix_sets) @@ -250,7 +246,7 @@ impl HashedPostState { tx: &TX, ) -> Result<(B256, TrieUpdates), StateRootError> { let sorted = self.clone().into_sorted(); - let prefix_sets = self.construct_prefix_sets(); + let prefix_sets = self.construct_prefix_sets().freeze(); StateRoot::from_tx(tx) .with_hashed_cursor_factory(HashedPostStateCursorFactory::new(tx, &sorted)) .with_prefix_sets(prefix_sets) From f0340018f7251ebb6129420381732e24486b4e4e Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Thu, 4 Jul 2024 16:13:06 +0200 Subject: [PATCH 2/2] rm unused import --- crates/trie/trie/src/state.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/trie/trie/src/state.rs b/crates/trie/trie/src/state.rs index cff265f09a99..f92b1dd08241 100644 --- a/crates/trie/trie/src/state.rs +++ b/crates/trie/trie/src/state.rs @@ -1,6 +1,6 @@ use crate::{ hashed_cursor::HashedPostStateCursorFactory, - prefix_set::{PrefixSetMut, TriePrefixSets, TriePrefixSetsMut}, + prefix_set::{PrefixSetMut, TriePrefixSetsMut}, updates::TrieUpdates, Nibbles, StateRoot, };