add state root from proofs bench and simplify account trie nodes coll…

…ection

crates/engine/tree/Cargo.toml

@@ -91,6 +91,10 @@ crossbeam-channel = "0.5.13"
 name = "channel_perf"
 harness = false
 
+[[bench]]
+name = "state_root_from_proofs"
+harness = false
+
 [features]
 test-utils = [
   "reth-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;
 

crates/engine/tree/src/tree/root.rs

@@ -304,7 +304,8 @@ where
     }
 }
 
-fn calculate_state_root_from_proofs<Factory>(
+/// Calculate state root from proofs.
+pub fn calculate_state_root_from_proofs<Factory>(
     view: ConsistentDbView<Factory>,
     input_nodes_sorted: &TrieUpdatesSorted,
     input_state_sorted: &HashedPostStateSorted,
@@ -397,14 +398,12 @@ where
             });
             let key = Nibbles::unpack(hashed_address);
             let proof = multiproof.account_subtree.iter().filter(|e| key.starts_with(e.0));
-            Ok(target_nodes(key.clone(), value, None, proof)?)
+            target_nodes(key.clone(), value, None, proof)
         })
-        .collect::<Vec<ProviderResult<BTreeMap<_, _>>>>()
-        .into_iter()
-        .collect::<Result<Vec<_>, _>>()?
-        .into_iter()
-        .flatten()
-        .collect::<BTreeMap<_, _>>();
+        .try_reduce(BTreeMap::new, |mut acc, map| {
+            acc.extend(map.into_iter());
+            Ok(acc)
+        })?;
 
     let state_root = next_root_from_proofs(account_trie_nodes, |key: Nibbles| {
         // Right pad the target with 0s.