feat(engine): proof fetching on state update for StateRootTask

crates/engine/tree/Cargo.toml

@@ -21,6 +21,7 @@ reth-consensus.workspace = true
 reth-engine-primitives.workspace = true
 reth-errors.workspace = true
 reth-evm.workspace = true
+reth-execution-errors.workspace = true
 reth-network-p2p.workspace = true
 reth-payload-builder-primitives.workspace = true
 reth-payload-builder.workspace = true
@@ -32,6 +33,7 @@ reth-prune.workspace = true
 reth-revm.workspace = true
 reth-stages-api.workspace = true
 reth-tasks.workspace = true
+reth-trie-db.workspace = true
 reth-trie-parallel.workspace = true
 reth-trie-sparse.workspace = true
 reth-trie.workspace = true
@@ -55,6 +57,7 @@ metrics.workspace = true
 reth-metrics = { workspace = true, features = ["common"] }
 
 # misc
+rayon.workspace = true
 tracing.workspace = true
 
 # optional deps for test-utils
@@ -77,6 +80,7 @@ reth-prune.workspace = true
 reth-rpc-types-compat.workspace = true
 reth-stages = { workspace = true, features = ["test-utils"] }
 reth-static-file.workspace = true
+reth-testing-utils.workspace = true
 reth-tracing.workspace = true
 
 # alloy
@@ -85,11 +89,16 @@ alloy-rlp.workspace = true
 assert_matches.workspace = true
 criterion.workspace = true
 crossbeam-channel = "0.5.13"
+rand.workspace = true
 
 [[bench]]
 name = "channel_perf"
 harness = false
 
+[[bench]]
+name = "state_root_from_proofs"
+harness = false
+
 [features]
 test-utils = [
     "reth-blockchain-tree/test-utils",
@@ -110,4 +119,6 @@ test-utils = [
     "reth-static-file",
     "reth-tracing",
     "reth-trie/test-utils",
+    "reth-prune-types?/test-utils",
+    "reth-trie-db/test-utils",
 ]

crates/engine/tree/benches/state_root_from_proofs.rs

@@ -0,0 +1,81 @@
+#![allow(missing_docs)]
+
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+use reth_engine_tree::tree::calculate_state_root_from_proofs;
+use reth_provider::{providers::ConsistentDbView, test_utils::create_test_provider_factory};
+use reth_trie::{
+    updates::TrieUpdatesSorted, HashedPostState, HashedPostStateSorted, HashedStorage, MultiProof,
+};
+use revm_primitives::{
+    keccak256, Account, AccountInfo, AccountStatus, Address, EvmStorage, EvmStorageSlot, HashMap,
+    HashSet, B256, U256,
+};
+
+fn create_test_state(size: usize) -> (HashMap<B256, HashSet<B256>>, HashedPostState) {
+    let mut state = HashedPostState::default();
+    let mut targets = HashMap::default();
+
+    for i in 0..size {
+        let address = Address::random();
+        let hashed_address = keccak256(address);
+
+        // Create account
+        let info = AccountInfo {
+            balance: U256::from(100 + i),
+            nonce: i as u64,
+            code_hash: B256::random(),
+            code: Default::default(),
+        };
+
+        // Create storage with multiple slots
+        let mut storage = EvmStorage::default();
+        let mut slots = HashSet::default();
+        for j in 0..100 {
+            let slot = U256::from(j);
+            let value = U256::from(100 + j);
+            storage.insert(slot, EvmStorageSlot::new(value));
+            slots.insert(keccak256(B256::from(slot)));
+        }
+
+        let account = Account { info, storage: storage.clone(), status: AccountStatus::Loaded };
+
+        state.accounts.insert(hashed_address, Some(account.info.into()));
+        state.storages.insert(
+            hashed_address,
+            HashedStorage::from_iter(
+                false,
+                storage.into_iter().map(|(k, v)| (keccak256(B256::from(k)), v.present_value)),
+            ),
+        );
+        targets.insert(hashed_address, slots);
+    }
+
+    (targets, state)
+}
+
+fn bench_state_root_collection(c: &mut Criterion) {
+    let factory = create_test_provider_factory();
+    let view = ConsistentDbView::new(factory, None);
+
+    let mut group = c.benchmark_group("state_root_collection");
+    for size in &[10, 100, 1000] {
+        let (_targets, state) = create_test_state(*size);
+        let multiproof = MultiProof::default();
+
+        group.bench_with_input(format!("size_{}", size), size, |b, _| {
+            b.iter(|| {
+                black_box(calculate_state_root_from_proofs(
+                    view.clone(),
+                    &TrieUpdatesSorted::default(),
+                    &HashedPostStateSorted::default(),
+                    multiproof.clone(),
+                    state.clone(),
+                ))
+            });
+        });
+    }
+    group.finish();
+}
+
+criterion_group!(benches, bench_state_root_collection);
+criterion_main!(benches);

crates/engine/tree/src/tree/mod.rs

@@ -76,6 +76,7 @@ pub use config::TreeConfig;
 pub use invalid_block_hook::{InvalidBlockHooks, NoopInvalidBlockHook};
 pub use persistence_state::PersistenceState;
 pub use reth_engine_primitives::InvalidBlockHook;
+pub use root::calculate_state_root_from_proofs;
 
 mod root;