chore(trie): return mutable prefix sets from HashedPostState::construct_prefix_sets (#9306)

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);

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)
                 },

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,

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,

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<B256, PrefixSetMut>,
+    /// A set of hashed addresses of destroyed accounts.
+    pub destroyed_accounts: HashSet<B256>,
+}
+
+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 {

crates/trie/trie/src/state.rs

@@ -1,6 +1,6 @@
 use crate::{
     hashed_cursor::HashedPostStateCursorFactory,
-    prefix_set::{PrefixSetMut, TriePrefixSets},
+    prefix_set::{PrefixSetMut, 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 {
+        TriePrefixSetsMut { account_prefix_set, storage_prefix_sets, destroyed_accounts }
-            account_prefix_set: account_prefix_set.freeze(),
-            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<TX: DbTx>(&self, tx: &TX) -> Result<B256, 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)
@@ -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)