diff --git a/crates/engine/tree/Cargo.toml b/crates/engine/tree/Cargo.toml index 5ca612e77c76..c1da351c7862 100644 --- a/crates/engine/tree/Cargo.toml +++ b/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", diff --git a/crates/engine/tree/benches/state_root_from_proofs.rs b/crates/engine/tree/benches/state_root_from_proofs.rs new file mode 100644 index 000000000000..4c8e85696ead --- /dev/null +++ b/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>, 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); diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs index 39843377684c..8f53ee4bd230 100644 --- a/crates/engine/tree/src/tree/mod.rs +++ b/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; diff --git a/crates/engine/tree/src/tree/root.rs b/crates/engine/tree/src/tree/root.rs index 54c192843242..08d7b3eb8a32 100644 --- a/crates/engine/tree/src/tree/root.rs +++ b/crates/engine/tree/src/tree/root.rs @@ -304,7 +304,8 @@ where } } -fn calculate_state_root_from_proofs( +/// Calculate state root from proofs. +pub fn calculate_state_root_from_proofs( view: ConsistentDbView, 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::>>>() - .into_iter() - .collect::, _>>()? - .into_iter() - .flatten() - .collect::>(); + .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.