perf(trie): parallelize HashedPostStateSorted::from_reverts

Parallelizes the hashing/sorting step using Rayon when account count
exceeds a threshold (2500). This alleviates CPU bottlenecks during large
state reverts or deep reorgs.

Closes #20049

Author: Andrurachi

crates/trie/db/Cargo.toml

@@ -24,6 +24,9 @@ alloy-primitives.workspace = true
 # tracing
 tracing.workspace = true
 
+# rayon
+rayon = { workspace = true, optional = true }
+
 [dev-dependencies]
 # reth
 reth-chainspec.workspace = true
@@ -48,7 +51,6 @@ serde_json.workspace = true
 similar-asserts.workspace = true
 
 criterion.workspace = true
-rayon.workspace = true
 rand.workspace = true
 
 [features]
@@ -62,6 +64,7 @@ serde = [
     "reth-primitives-traits/serde",
     "revm-database/serde",
     "revm/serde",
+    "rand/serde",
 ]
 test-utils = [
     "reth-trie-common/test-utils",
@@ -73,6 +76,13 @@ test-utils = [
     "reth-trie/test-utils",
 ]
 
+parallel-from-reverts = ["dep:rayon"]
+
 [[bench]]
 name = "sorting_par_exp"
-harness = false
+harness = false
+required-features = ["parallel-from-reverts"]
+
+[[bench]]
+name = "integration_bench"
+harness = false

crates/trie/db/benches/integration_bench.rs

@@ -0,0 +1,110 @@
+#![allow(missing_docs, unreachable_pub)]
+
+use alloy_primitives::{Address, B256, U256};
+use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
+use rand::{rngs::StdRng, Rng, SeedableRng};
+use reth_db::{
+    test_utils::create_test_rw_db,
+    transaction::{DbTx, DbTxMut},
+    Database, DatabaseEnv,
+};
+use reth_db_api::{
+    models::{AccountBeforeTx, BlockNumberAddress},
+    tables,
+};
+use reth_primitives_traits::{Account, StorageEntry};
+use reth_trie::{HashedPostStateSorted, KeccakKeyHasher};
+use reth_trie_db::DatabaseHashedPostState;
+use std::sync::Arc;
+
+// Populate a temporary database with synthetic revert data.
+type TestDb = Arc<reth_db::test_utils::TempDatabase<DatabaseEnv>>;
+
+#[derive(Clone, Copy, Debug)]
+enum Distribution {
+    Uniform(usize),
+    Skewed { light: usize, heavy: usize },
+}
+
+fn seed_db(num_accounts: usize, dist: Distribution) -> TestDb {
+    let db = create_test_rw_db();
+    let tx = db.tx_mut().expect("failed to create rw tx");
+
+    let mut rng = StdRng::seed_from_u64(12345);
+    let block = 1;
+
+    for idx in 0..num_accounts {
+        let address = Address::random();
+
+        let account =
+            Account { nonce: idx as u64, balance: U256::from(idx as u64), bytecode_hash: None };
+
+        tx.put::<tables::AccountChangeSets>(
+            block,
+            AccountBeforeTx { address, info: Some(account) },
+        )
+        .expect("failed to insert account changeset");
+
+        let slots_count = match dist {
+            Distribution::Uniform(n) => n,
+            Distribution::Skewed { light, heavy } => {
+                if rng.random_bool(0.05) {
+                    heavy
+                } else {
+                    light
+                }
+            }
+        };
+
+        for slot_idx in 0..slots_count {
+            let slot_key = B256::random();
+            let value = U256::from(slot_idx);
+
+            tx.put::<tables::StorageChangeSets>(
+                BlockNumberAddress((block, address)),
+                StorageEntry { key: slot_key, value },
+            )
+            .expect("failed to insert storage changeset");
+        }
+    }
+
+    tx.commit().expect("failed to commit seeded db");
+    db
+}
+
+fn bench_db_execution(c: &mut Criterion) {
+    let mut group = c.benchmark_group("Integration_DB_Sort");
+    group.measurement_time(std::time::Duration::from_secs(10));
+
+    // Scenarios to test:
+    // - Below Threshold: Should be identical (Sequential)
+    // - Above Threshold: Should show some speedup with feature flag
+    // - Skewed: Should show speedup with feature flag
+    let scenarios = vec![
+        ("Small_Under_Threshold", 1_000, Distribution::Uniform(50)),
+        ("Large_Uniform", 10_000, Distribution::Uniform(20)),
+        ("Large_Skewed", 10_000, Distribution::Skewed { light: 4, heavy: 2_000 }),
+    ];
+
+    for (name, accounts, dist) in scenarios {
+        let db = seed_db(accounts, dist);
+        let range = 1..=1;
+
+        group.bench_function(BenchmarkId::new("from_reverts", name), |b| {
+            b.iter(|| {
+                let tx = db.tx().expect("failed to create ro tx");
+
+                let state =
+                    HashedPostStateSorted::from_reverts::<KeccakKeyHasher>(&tx, range.clone())
+                        .expect("failed to calculate state");
+
+                black_box(state);
+            });
+        });
+    }
+
+    group.finish();
+}
+
+criterion_group!(benches, bench_db_execution);
+criterion_main!(benches);

crates/trie/db/src/state.rs

@@ -19,6 +19,9 @@ use std::{
 };
 use tracing::{debug, instrument};
 
+#[cfg(feature = "parallel-from-reverts")]
+use rayon::iter::{IntoParallelIterator, ParallelIterator};
+
 /// Extends [`StateRoot`] with operations specific for working with a database transaction.
 pub trait DatabaseStateRoot<'a, TX>: Sized {
     /// Create a new [`StateRoot`] instance.
@@ -280,14 +283,37 @@ impl<TX: DbTx> DatabaseHashedPostState<TX> for HashedPostStateSorted {
             }
         }
 
-        // Sort storage slots and convert to HashedStorageSorted
-        let hashed_storages = storages
-            .into_iter()
-            .map(|(address, mut slots)| {
-                slots.sort_unstable_by_key(|(slot, _)| *slot);
-                (address, HashedStorageSorted { storage_slots: slots, wiped: false })
-            })
-            .collect();
+        // Threshold based on benchmark
+        const PARALLEL_THRESHOLD: usize = 2_500;
+
+        // Check Feature Flag AND Threshold
+        let use_parallel =
+            cfg!(feature = "parallel-from-reverts") && storages.len() >= PARALLEL_THRESHOLD;
+
+        let hashed_storages = if use_parallel {
+            #[cfg(feature = "parallel-from-reverts")]
+            {
+                storages
+                    .into_par_iter()
+                    .map(|(address, mut slots)| {
+                        slots.sort_unstable_by_key(|(slot, _)| *slot);
+                        (address, HashedStorageSorted { storage_slots: slots, wiped: false })
+                    })
+                    .collect()
+            }
+            #[cfg(not(feature = "parallel-from-reverts"))]
+            {
+                unreachable!()
+            }
+        } else {
+            storages
+                .into_iter()
+                .map(|(address, mut slots)| {
+                    slots.sort_unstable_by_key(|(slot, _)| *slot);
+                    (address, HashedStorageSorted { storage_slots: slots, wiped: false })
+                })
+                .collect()
+        };
 
         Ok(Self::new(accounts, hashed_storages))
     }