From 5bb89ed7f9c4402b634c095880d011476c1b1d72 Mon Sep 17 00:00:00 2001
From: Roman Krasiuk <rokrassyuk@gmail.com>
Date: Mon, 15 Jul 2024 15:20:34 +0200
Subject: [PATCH 1/2] feat(trie): sorted iterators for updated hashed state
 entries

---
 Cargo.lock                    |  1 +
 crates/trie/trie/Cargo.toml   |  1 +
 crates/trie/trie/src/state.rs | 42 ++++++++++++++++++++++++++++++++++-
 3 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index 52184d98f177..4e8762e07102 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -8639,6 +8639,7 @@ dependencies = [
  "auto_impl",
  "criterion",
  "derive_more",
+ "itertools 0.13.0",
  "metrics",
  "once_cell",
  "proptest",
diff --git a/crates/trie/trie/Cargo.toml b/crates/trie/trie/Cargo.toml
index 04b03014e33f..8c9ddd43bd03 100644
--- a/crates/trie/trie/Cargo.toml
+++ b/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 }
diff --git a/crates/trie/trie/src/state.rs b/crates/trie/trie/src/state.rs
index c6c93c0b3620..c93fd47da5b1 100644
--- a/crates/trie/trie/src/state.rs
+++ b/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,6 +351,17 @@ 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 {
@@ -345,10 +369,26 @@ pub struct HashedStorageSorted {
     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::*;

From 272fc9bffc25622634dfb041d5a81769916da657 Mon Sep 17 00:00:00 2001
From: Roman Krasiuk <rokrassyuk@gmail.com>
Date: Mon, 15 Jul 2024 15:28:17 +0200
Subject: [PATCH 2/2] const

---
 crates/trie/trie/src/state.rs | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/crates/trie/trie/src/state.rs b/crates/trie/trie/src/state.rs
index c93fd47da5b1..187cc842b0b9 100644
--- a/crates/trie/trie/src/state.rs
+++ b/crates/trie/trie/src/state.rs
@@ -332,12 +332,12 @@ pub struct HashedPostStateSorted {
 
 impl HashedPostStateSorted {
     /// Returns reference to hashed accounts.
-    pub fn accounts(&self) -> &HashedAccountsSorted {
+    pub const fn accounts(&self) -> &HashedAccountsSorted {
         &self.accounts
     }
 
     /// Returns reference to hashed account storages.
-    pub fn account_storages(&self) -> &HashMap<B256, HashedStorageSorted> {
+    pub const fn account_storages(&self) -> &HashMap<B256, HashedStorageSorted> {
         &self.storages
     }
 }
@@ -353,7 +353,7 @@ pub struct HashedAccountsSorted {
 
 impl HashedAccountsSorted {
     /// Returns a sorted iterator over updated accounts.
-    pub fn accounts_sorted(self) -> impl Iterator<Item = (B256, Option<Account>)> {
+    pub fn accounts_sorted(&self) -> impl Iterator<Item = (B256, Option<Account>)> {
         self.accounts
             .iter()
             .map(|(address, account)| (*address, Some(*account)))
@@ -375,7 +375,7 @@ pub struct HashedStorageSorted {
 
 impl HashedStorageSorted {
     /// Returns `true` if the account was wiped.
-    pub fn is_wiped(&self) -> bool {
+    pub const fn is_wiped(&self) -> bool {
         self.wiped
     }