Moved HashedPostState to trie-common crate (#14230)

Co-authored-by: DarkLord017 <sambhavjain170944@gmail.com>
Co-authored-by: Matthias Seitz <matthias.seitz@outlook.de>

crates/trie/common/Cargo.toml

@@ -28,6 +28,9 @@ derive_more.workspace = true
 itertools = { workspace = true, features = ["use_alloc"] }
 nybbles = { workspace = true, features = ["rlp"] }
 
+# reth
+revm.workspace = true
+
 # `serde` feature
 serde = { workspace = true, optional = true }
 
@@ -38,6 +41,9 @@ hash-db = { version = "=0.15.2", optional = true }
 plain_hasher = { version = "0.2", optional = true }
 arbitrary = { workspace = true, features = ["derive"], optional = true }
 
+# misc
+rayon = { workspace = true, optional = true }
+
 [dev-dependencies]
 reth-primitives-traits = { workspace = true, features = ["serde"] }
 reth-codecs.workspace = true
@@ -74,6 +80,7 @@ std = [
     "serde?/std",
     "serde_with?/std",
     "serde_json/std",
+    "revm/std",
 ]
 eip1186 = [
     "alloy-rpc-types-eth/serde",
@@ -89,6 +96,7 @@ serde = [
     "alloy-rpc-types-eth?/serde",
     "reth-primitives-traits/serde",
     "reth-codecs?/serde",
+    "revm/serde",
 ]
 reth-codec = [
     "dep:reth-codecs",
@@ -106,6 +114,7 @@ test-utils = [
     "arbitrary",
     "reth-primitives-traits/test-utils",
     "reth-codecs/test-utils",
+    "revm/test-utils",
 ]
 arbitrary = [
     "std",
@@ -119,7 +128,9 @@ arbitrary = [
     "nybbles/arbitrary",
     "reth-codecs/arbitrary",
     "alloy-rpc-types-eth?/arbitrary",
+    "revm/arbitrary",
 ]
+rayon = ["dep:rayon"]
 
 [[bench]]
 name = "prefix_set"

crates/trie/common/src/hashedstate.rs

@@ -1,18 +1,22 @@
 use crate::{
     prefix_set::{PrefixSetMut, TriePrefixSetsMut},
-    Nibbles,
+    KeyHasher, Nibbles,
 };
+use alloc::{borrow::Cow, vec::Vec};
 use alloy_primitives::{
     keccak256,
     map::{hash_map, B256HashMap, B256HashSet, HashMap, HashSet},
     Address, B256, U256,
 };
 use itertools::Itertools;
-use rayon::prelude::{IntoParallelIterator, ParallelIterator};
 use reth_primitives_traits::Account;
-use reth_trie_common::KeyHasher;
 use revm::db::{AccountStatus, BundleAccount};
-use std::borrow::Cow;
+
+#[cfg(feature = "rayon")]
+pub use rayon::*;
+
+#[cfg(feature = "rayon")]
+use rayon::prelude::{IntoParallelIterator, ParallelIterator};
 
 /// Representation of in-memory hashed state.
 #[derive(PartialEq, Eq, Clone, Default, Debug)]
@@ -28,6 +32,7 @@ impl HashedPostState {
     /// Hashes all changed accounts and storage entries that are currently stored in the bundle
     /// state.
     #[inline]
+    #[cfg(feature = "rayon")]
     pub fn from_bundle_state<'a, KH: KeyHasher>(
         state: impl IntoParallelIterator<Item = (&'a Address, &'a BundleAccount)>,
     ) -> Self {
@@ -55,6 +60,37 @@ impl HashedPostState {
         Self { accounts, storages }
     }
 
+    /// Initialize [`HashedPostState`] from bundle state.
+    /// Hashes all changed accounts and storage entries that are currently stored in the bundle
+    /// state.
+    #[cfg(not(feature = "rayon"))]
+    pub fn from_bundle_state<'a, KH: KeyHasher>(
+        state: impl IntoIterator<Item = (&'a Address, &'a BundleAccount)>,
+    ) -> Self {
+        let hashed = state
+            .into_iter()
+            .map(|(address, account)| {
+                let hashed_address = KH::hash_key(address);
+                let hashed_account = account.info.as_ref().map(Into::into);
+                let hashed_storage = HashedStorage::from_plain_storage(
+                    account.status,
+                    account.storage.iter().map(|(slot, value)| (slot, &value.present_value)),
+                );
+                (hashed_address, (hashed_account, hashed_storage))
+            })
+            .collect::<Vec<(B256, (Option<Account>, HashedStorage))>>();
+
+        let mut accounts = HashMap::with_capacity_and_hasher(hashed.len(), Default::default());
+        let mut storages = HashMap::with_capacity_and_hasher(hashed.len(), Default::default());
+        for (address, (account, storage)) in hashed {
+            accounts.insert(address, account);
+            if !storage.is_empty() {
+                storages.insert(address, storage);
+            }
+        }
+        Self { accounts, storages }
+    }
+
     /// Construct [`HashedPostState`] from a single [`HashedStorage`].
     pub fn from_hashed_storage(hashed_address: B256, storage: HashedStorage) -> Self {
         Self {
@@ -260,9 +296,9 @@ impl HashedStorage {
 #[derive(PartialEq, Eq, Clone, Default, Debug)]
 pub struct HashedPostStateSorted {
     /// Updated state of accounts.
-    pub(crate) accounts: HashedAccountsSorted,
+    pub accounts: HashedAccountsSorted,
     /// Map of hashed addresses to hashed storage.
-    pub(crate) storages: B256HashMap<HashedStorageSorted>,
+    pub storages: B256HashMap<HashedStorageSorted>,
 }
 
 impl HashedPostStateSorted {
@@ -289,9 +325,9 @@ impl HashedPostStateSorted {
 #[derive(Clone, Eq, PartialEq, Default, Debug)]
 pub struct HashedAccountsSorted {
     /// Sorted collection of hashed addresses and their account info.
-    pub(crate) accounts: Vec<(B256, Account)>,
+    pub accounts: Vec<(B256, Account)>,
     /// Set of destroyed account keys.
-    pub(crate) destroyed_accounts: B256HashSet,
+    pub destroyed_accounts: B256HashSet,
 }
 
 impl HashedAccountsSorted {
@@ -309,11 +345,11 @@ impl HashedAccountsSorted {
 #[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)>,
+    pub non_zero_valued_slots: Vec<(B256, U256)>,
     /// Slots that have been zero valued.
-    pub(crate) zero_valued_slots: B256HashSet,
+    pub zero_valued_slots: B256HashSet,
     /// Flag indicating whether the storage was wiped or not.
-    pub(crate) wiped: bool,
+    pub wiped: bool,
 }
 
 impl HashedStorageSorted {
@@ -335,8 +371,8 @@ impl HashedStorageSorted {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::KeccakKeyHasher;
     use alloy_primitives::Bytes;
-    use reth_trie_common::KeccakKeyHasher;
     use revm::{
         db::{states::StorageSlot, StorageWithOriginalValues},
         primitives::{AccountInfo, Bytecode},

crates/trie/common/src/lib.rs

@@ -11,6 +11,10 @@
 
 extern crate alloc;
 
+/// In-memory hashed state.
+mod hashedstate;
+pub use hashedstate::*;
+
 /// The implementation of hash builder.
 pub mod hash_builder;
 

crates/trie/trie/Cargo.toml

@@ -18,7 +18,7 @@ reth-primitives-traits.workspace = true
 reth-stages-types.workspace = true
 reth-storage-errors.workspace = true
 reth-trie-sparse.workspace = true
-reth-trie-common.workspace = true
+reth-trie-common = { workspace = true, features = ["rayon"] }
 
 revm.workspace = true
 
@@ -33,7 +33,6 @@ alloy-trie.workspace = true
 tracing.workspace = true
 
 # misc
-rayon.workspace = true
 auto_impl.workspace = true
 itertools.workspace = true
 

crates/trie/trie/src/hashed_cursor/post_state.rs

@@ -1,11 +1,9 @@
 use super::{HashedCursor, HashedCursorFactory, HashedStorageCursor};
-use crate::{
-    forward_cursor::ForwardInMemoryCursor, HashedAccountsSorted, HashedPostStateSorted,
-    HashedStorageSorted,
-};
+use crate::forward_cursor::ForwardInMemoryCursor;
 use alloy_primitives::{map::B256HashSet, B256, U256};
 use reth_primitives_traits::Account;
 use reth_storage_errors::db::DatabaseError;
+use reth_trie_common::{HashedAccountsSorted, HashedPostStateSorted, HashedStorageSorted};
 
 /// The hashed cursor factory for the post state.
 #[derive(Clone, Debug)]

crates/trie/trie/src/lib.rs

@@ -4,6 +4,7 @@
 //!
 //! ## Feature Flags
 //!
+//! - `rayon`: uses rayon for parallel [`HashedPostState`] creation.
 //! - `test-utils`: Export utilities for testing
 
 #![doc(
@@ -28,10 +29,6 @@ pub mod walker;
 /// The iterators for traversing existing intermediate hashes and updated trie leaves.
 pub mod node_iter;
 
-/// In-memory hashed state.
-mod state;
-pub use state::*;
-
 /// Input for trie computation.
 mod input;
 pub use input::TrieInput;

crates/trie/trie/src/witness.rs

@@ -3,8 +3,9 @@ use crate::{
     prefix_set::TriePrefixSetsMut,
     proof::{Proof, ProofBlindedProviderFactory},
     trie_cursor::TrieCursorFactory,
-    HashedPostState,
 };
+use reth_trie_common::HashedPostState;
+
 use alloy_primitives::{
     keccak256,
     map::{B256HashMap, B256HashSet, Entry, HashMap},