feat(trie): sorted iterators for updated hashed state entries

crates/trie/trie/Cargo.toml

@@ -32,6 +32,7 @@ tracing.workspace = true
 rayon.workspace = true
 derive_more.workspace = true
 auto_impl.workspace = true
+itertools.workspace = true
 
 # `metrics` feature
 reth-metrics = { workspace = true, optional = true }

crates/trie/trie/src/state.rs

@@ -5,6 +5,7 @@ use crate::{
     updates::TrieUpdates,
     Nibbles, StateRoot,
 };
+use itertools::Itertools;
 use rayon::prelude::{IntoParallelIterator, ParallelIterator};
 use reth_db::{tables, DatabaseError};
 use reth_db_api::{
@@ -329,6 +330,18 @@ pub struct HashedPostStateSorted {
     pub(crate) storages: HashMap<B256, HashedStorageSorted>,
 }
 
+impl HashedPostStateSorted {
+    /// Returns reference to hashed accounts.
+    pub fn accounts(&self) -> &HashedAccountsSorted {
+        &self.accounts
+    }
+
+    /// Returns reference to hashed account storages.
+    pub fn account_storages(&self) -> &HashMap<B256, HashedStorageSorted> {
+        &self.storages
+    }
+}
+
 /// Sorted account state optimized for iterating during state trie calculation.
 #[derive(Clone, Eq, PartialEq, Debug)]
 pub struct HashedAccountsSorted {
@@ -338,17 +351,44 @@ pub struct HashedAccountsSorted {
     pub(crate) destroyed_accounts: HashSet<B256>,
 }
 
+impl HashedAccountsSorted {
+    /// Returns a sorted iterator over updated accounts.
+    pub fn accounts_sorted(self) -> impl Iterator<Item = (B256, Option<Account>)> {
+        self.accounts
+            .iter()
+            .map(|(address, account)| (*address, Some(*account)))
+            .chain(self.destroyed_accounts.iter().map(|address| (*address, None)))
+            .sorted_by_key(|entry| *entry.0)
+    }
+}
+
 /// Sorted hashed storage optimized for iterating during state trie calculation.
 #[derive(Clone, Eq, PartialEq, Debug)]
 pub struct HashedStorageSorted {
     /// Sorted hashed storage slots with non-zero value.
     pub(crate) non_zero_valued_slots: Vec<(B256, U256)>,
     /// Slots that have been zero valued.
     pub(crate) zero_valued_slots: HashSet<B256>,
-    /// Flag indicating hether the storage was wiped or not.
+    /// Flag indicating whether the storage was wiped or not.
     pub(crate) wiped: bool,
 }
 
+impl HashedStorageSorted {
+    /// Returns `true` if the account was wiped.
+    pub fn is_wiped(&self) -> bool {
+        self.wiped
+    }
+
+    /// Returns a sorted iterator over updated storage slots.
+    pub fn storage_slots_sorted(&self) -> impl Iterator<Item = (B256, U256)> {
+        self.non_zero_valued_slots
+            .iter()
+            .map(|(hashed_slot, value)| (*hashed_slot, *value))
+            .chain(self.zero_valued_slots.iter().map(|hashed_slot| (*hashed_slot, U256::ZERO)))
+            .sorted_by_key(|entry| *entry.0)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;