feat(trie): hash post state in parallel (#7190)

Signed-off-by: jsvisa <delweng@gmail.com>
Co-authored-by: Roman Krasiuk <rokrassyuk@gmail.com>

crates/trie/trie/Cargo.toml

@@ -26,6 +26,7 @@ alloy-rlp.workspace = true
 tracing.workspace = true
 
 # misc
+rayon.workspace = true
 derive_more.workspace = true
 auto_impl.workspace = true
 
@@ -62,3 +63,7 @@ test-utils = ["triehash"]
 [[bench]]
 name = "prefix_set"
 harness = false
+
+[[bench]]
+name = "hash_post_state"
+harness = false

crates/trie/trie/benches/hash_post_state.rs

@@ -0,0 +1,80 @@
+#![allow(missing_docs, unreachable_pub)]
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use proptest::{prelude::*, strategy::ValueTree, test_runner::TestRunner};
+use reth_primitives::{keccak256, revm::compat::into_reth_acc, Address, B256, U256};
+use reth_trie::{HashedPostState, HashedStorage};
+use revm::db::{states::BundleBuilder, BundleAccount};
+use std::collections::HashMap;
+
+pub fn hash_post_state(c: &mut Criterion) {
+    let mut group = c.benchmark_group("Hash Post State");
+    group.sample_size(20);
+
+    for size in [100, 1_000, 3_000, 5_000, 10_000] {
+        let state = generate_test_data(size);
+
+        // sequence
+        group.bench_function(BenchmarkId::new("sequence hashing", size), |b| {
+            b.iter(|| from_bundle_state_seq(&state))
+        });
+
+        // parallel
+        group.bench_function(BenchmarkId::new("parallel hashing", size), |b| {
+            b.iter(|| HashedPostState::from_bundle_state(&state))
+        });
+    }
+}
+
+fn from_bundle_state_seq(state: &HashMap<Address, BundleAccount>) -> HashedPostState {
+    let mut this = HashedPostState::default();
+
+    for (address, account) in state {
+        let hashed_address = keccak256(address);
+        this.accounts.insert(hashed_address, account.info.clone().map(into_reth_acc));
+
+        let hashed_storage = HashedStorage::from_iter(
+            account.status.was_destroyed(),
+            account
+                .storage
+                .iter()
+                .map(|(key, value)| (keccak256(B256::new(key.to_be_bytes())), value.present_value)),
+        );
+        this.storages.insert(hashed_address, hashed_storage);
+    }
+    this
+}
+
+fn generate_test_data(size: usize) -> HashMap<Address, BundleAccount> {
+    let storage_size = 1_000;
+    let mut runner = TestRunner::new(ProptestConfig::default());
+
+    use proptest::collection::hash_map;
+    let state = hash_map(
+        any::<Address>(),
+        hash_map(
+            any::<U256>(), // slot
+            (
+                any::<U256>(), // old value
+                any::<U256>(), // new value
+            ),
+            storage_size,
+        ),
+        size,
+    )
+    .new_tree(&mut runner)
+    .unwrap()
+    .current();
+
+    let mut bundle_builder = BundleBuilder::default();
+
+    for (address, storage) in state.into_iter() {
+        bundle_builder = bundle_builder.state_storage(address, storage);
+    }
+
+    let bundle_state = bundle_builder.build();
+
+    bundle_state.state
+}
+
+criterion_group!(post_state, hash_post_state);
+criterion_main!(post_state);

crates/trie/trie/src/state.rs

@@ -4,6 +4,7 @@ use crate::{
     updates::TrieUpdates,
     StateRoot,
 };
+use rayon::prelude::{IntoParallelIterator, ParallelIterator};
 use reth_db::{
     cursor::DbCursorRO,
     models::{AccountBeforeTx, BlockNumberAddress},
@@ -36,22 +37,30 @@ impl HashedPostState {
     /// Hashes all changed accounts and storage entries that are currently stored in the bundle
     /// state.
     pub fn from_bundle_state<'a>(
-        state: impl IntoIterator<Item = (&'a Address, &'a BundleAccount)>,
+        state: impl IntoParallelIterator<Item = (&'a Address, &'a BundleAccount)>,
     ) -> Self {
-        let mut this = Self::default();
-        for (address, account) in state {
-            let hashed_address = keccak256(address);
-            this.accounts.insert(hashed_address, account.info.clone().map(into_reth_acc));
+        let hashed = state
+            .into_par_iter()
+            .map(|(address, account)| {
+                let hashed_address = keccak256(address);
+                let hashed_account = account.info.clone().map(into_reth_acc);
+                let hashed_storage = HashedStorage::from_iter(
+                    account.status.was_destroyed(),
+                    account.storage.iter().map(|(key, value)| {
+                        (keccak256(B256::new(key.to_be_bytes())), value.present_value)
+                    }),
+                );
+                (hashed_address, (hashed_account, hashed_storage))
+            })
+            .collect::<Vec<(B256, (Option<Account>, HashedStorage))>>();
 
-            let hashed_storage = HashedStorage::from_iter(
-                account.status.was_destroyed(),
-                account.storage.iter().map(|(key, value)| {
-                    (keccak256(B256::new(key.to_be_bytes())), value.present_value)
-                }),
-            );
-            this.storages.insert(hashed_address, hashed_storage);
+        let mut accounts = HashMap::with_capacity(hashed.len());
+        let mut storages = HashMap::with_capacity(hashed.len());
+        for (address, (account, storage)) in hashed {
+            accounts.insert(address, account);
+            storages.insert(address, storage);
         }
-        this
+        Self { accounts, storages }
     }
 
     /// Initialize [`HashedPostState`] from revert range.
@@ -325,6 +334,12 @@ pub struct HashedStorageSorted {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use reth_db::{database::Database, test_utils::create_test_rw_db};
+    use reth_primitives::hex;
+    use revm::{
+        db::states::BundleState,
+        primitives::{AccountInfo, HashMap},
+    };
 
     #[test]
     fn hashed_state_wiped_extension() {
@@ -399,4 +414,34 @@ mod tests {
         );
         assert_eq!(account_storage.map(|st| st.wiped), Some(true));
     }
+
+    #[test]
+    fn from_bundle_state_with_rayon() {
+        let address1 = Address::with_last_byte(1);
+        let address2 = Address::with_last_byte(2);
+        let slot1 = U256::from(1015);
+        let slot2 = U256::from(2015);
+
+        let account1 = AccountInfo { nonce: 1, ..Default::default() };
+        let account2 = AccountInfo { nonce: 2, ..Default::default() };
+
+        let bundle_state = BundleState::builder(2..=2)
+            .state_present_account_info(address1, account1)
+            .state_present_account_info(address2, account2)
+            .state_storage(address1, HashMap::from([(slot1, (U256::ZERO, U256::from(10)))]))
+            .state_storage(address2, HashMap::from([(slot2, (U256::ZERO, U256::from(20)))]))
+            .build();
+        assert_eq!(bundle_state.reverts.len(), 1);
+
+        let post_state = HashedPostState::from_bundle_state(&bundle_state.state);
+        assert_eq!(post_state.accounts.len(), 2);
+        assert_eq!(post_state.storages.len(), 2);
+
+        let db = create_test_rw_db();
+        let tx = db.tx().expect("failed to create transaction");
+        assert_eq!(
+            post_state.state_root(&tx).unwrap(),
+            hex!("b464525710cafcf5d4044ac85b72c08b1e76231b8d91f288fe438cc41d8eaafd")
+        );
+    }
 }