diff --git a/Cargo.lock b/Cargo.lock index e9b475ecb004..b0350a6bcbc6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6527,6 +6527,7 @@ dependencies = [ "alloy-primitives", "alloy-signer", "alloy-signer-local", + "codspeed-criterion-compat", "derive_more", "metrics", "parking_lot", diff --git a/crates/chain-state/Cargo.toml b/crates/chain-state/Cargo.toml index 639b211d54ec..ac77c6fd9fae 100644 --- a/crates/chain-state/Cargo.toml +++ b/crates/chain-state/Cargo.toml @@ -52,6 +52,7 @@ reth-testing-utils.workspace = true alloy-signer.workspace = true alloy-signer-local.workspace = true alloy-consensus.workspace = true +criterion.workspace = true rand.workspace = true [features] @@ -65,3 +66,7 @@ test-utils = [ "reth-trie/test-utils", "revm/test-utils", ] + +[[bench]] +name = "trie_state" +harness = false diff --git a/crates/chain-state/benches/trie_state.rs b/crates/chain-state/benches/trie_state.rs new file mode 100644 index 000000000000..1466abf5cd53 --- /dev/null +++ b/crates/chain-state/benches/trie_state.rs @@ -0,0 +1,296 @@ +//! Usage: `cargo bench --package reth-chain-state --bench trie_state` + +#![allow(missing_docs)] + +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; + +use alloy_primitives::{Address, B256, U256}; +use criterion::{black_box, criterion_group, criterion_main, Criterion}; +use rand::{seq::SliceRandom, Rng}; +use reth_chain_state::{ + ExecutedBlock, ExecutedBlockWithTrieUpdates, MemoryOverlayStateProviderRef, +}; +use reth_primitives::{Account, NodePrimitives}; +use reth_storage_api::{noop::NoopProvider, StorageRootProvider}; +use reth_trie::{ + updates::{StorageTrieUpdates, TrieUpdates}, + BranchNodeCompact, HashedPostState, HashedStorage, Nibbles, +}; + +struct HashedPostStateFactory { + num_updated_accounts: usize, + num_removed_accounts: usize, + hashed_address_choices: Vec, + num_storages: usize, + num_changes_per_storage: usize, + storage_key_choices: Vec, +} + +impl HashedPostStateFactory { + fn erc20() -> Self { + Self { + num_updated_accounts: 30000, + num_removed_accounts: 0, + hashed_address_choices: (0..100000).map(|_| B256::random()).collect(), + num_storages: 30000, + num_changes_per_storage: 1, + storage_key_choices: vec![B256::random()], + } + } + + fn raw_transfer() -> Self { + Self { + num_updated_accounts: 75000, + num_removed_accounts: 0, + hashed_address_choices: (0..100000).map(|_| B256::random()).collect(), + num_storages: 75000, + num_changes_per_storage: 0, + storage_key_choices: Vec::new(), + } + } + + fn uniswap() -> Self { + Self { + num_updated_accounts: 6360, + num_removed_accounts: 0, + hashed_address_choices: (0..100000).map(|_| B256::random()).collect(), + num_storages: 6360, + num_changes_per_storage: 3, + storage_key_choices: vec![B256::random(), B256::random(), B256::random()], + } + } + + fn generate(&self, rng: &mut R) -> HashedPostState { + HashedPostState { + accounts: self + .hashed_address_choices + .choose_multiple(rng, self.num_updated_accounts + self.num_removed_accounts) + .enumerate() + .map(|(i, k)| (*k, (i < self.num_updated_accounts).then_some(Account::default()))) + .collect(), + storages: self + .hashed_address_choices + .choose_multiple(rng, self.num_storages) + .map(|k| { + ( + *k, + HashedStorage { + wiped: false, + storage: self + .storage_key_choices + .choose_multiple(rng, self.num_changes_per_storage) + .map(|k| (*k, U256::ZERO)) + .collect(), + }, + ) + }) + .collect(), + } + } +} + +struct TrieUpdatesFactory { + num_updated_nodes: usize, + num_removed_nodes: usize, + account_nibbles_choices: Vec, + num_storage_tries: usize, + hashed_address_choices: Vec, + storage_nibbles_choices: Vec, + num_updated_nodes_per_storage_trie: usize, + num_removed_nodes_per_storage_trie: usize, +} + +impl TrieUpdatesFactory { + fn erc20() -> Self { + Self { + num_updated_nodes: 6600, + num_removed_nodes: 0, + account_nibbles_choices: (0..7600) + .map(|_| Nibbles::unpack(B256::random().as_slice())) + .collect(), + num_storage_tries: 70000, + hashed_address_choices: (0..100000).map(|_| B256::random()).collect(), + storage_nibbles_choices: vec![Nibbles::unpack(B256::random().as_slice())], + num_updated_nodes_per_storage_trie: 1, + num_removed_nodes_per_storage_trie: 0, + } + } + + fn raw_transfer() -> Self { + Self { + num_updated_nodes: 7500, + num_removed_nodes: 0, + account_nibbles_choices: (0..7600) + .map(|_| Nibbles::unpack(B256::random().as_slice())) + .collect(), + num_storage_tries: 100000, + hashed_address_choices: (0..100000).map(|_| B256::random()).collect(), + storage_nibbles_choices: Vec::new(), + num_updated_nodes_per_storage_trie: 0, + num_removed_nodes_per_storage_trie: 0, + } + } + + fn uniswap() -> Self { + Self { + num_updated_nodes: 4200, + num_removed_nodes: 0, + account_nibbles_choices: (0..6000) + .map(|_| Nibbles::unpack(B256::random().as_slice())) + .collect(), + num_storage_tries: 32000, + hashed_address_choices: (0..60000).map(|_| B256::random()).collect(), + storage_nibbles_choices: vec![Nibbles::unpack(B256::random().as_slice())], + num_updated_nodes_per_storage_trie: 1, + num_removed_nodes_per_storage_trie: 0, + } + } + + fn generate(&self, rng: &mut R) -> TrieUpdates { + TrieUpdates { + changed_nodes: self + .account_nibbles_choices + .choose_multiple(rng, self.num_updated_nodes + self.num_removed_nodes) + .enumerate() + .map(|(i, k)| { + ( + k.clone(), + (i < self.num_updated_nodes).then_some(BranchNodeCompact::default()), + ) + }) + .collect(), + storage_tries: self + .hashed_address_choices + .choose_multiple(rng, self.num_storage_tries) + .map(|k| { + ( + *k, + StorageTrieUpdates { + is_deleted: false, + changed_nodes: (self + .storage_nibbles_choices + .choose_multiple( + rng, + self.num_updated_nodes_per_storage_trie + + self.num_removed_nodes_per_storage_trie, + ) + .enumerate() + .map(|(i, k)| { + ( + k.clone(), + (i < self.num_updated_nodes_per_storage_trie) + .then_some(BranchNodeCompact::default()), + ) + }) + .collect()), + }, + ) + }) + .collect(), + } + } +} + +fn generate_erc20_blocks(len: usize) -> Vec> { + let mut rng = rand::thread_rng(); + let hashed_post_state_factory = HashedPostStateFactory::erc20(); + let trie_updates_factory = TrieUpdatesFactory::erc20(); + (0..len) + .map(|_| ExecutedBlockWithTrieUpdates { + block: ExecutedBlock { + recovered_block: Arc::default(), + execution_output: Arc::default(), + hashed_state: Arc::new(hashed_post_state_factory.generate(&mut rng)), + }, + trie: Arc::new(trie_updates_factory.generate(&mut rng)), + }) + .collect() +} + +fn generate_raw_transfer_blocks( + len: usize, +) -> Vec> { + let mut rng = rand::thread_rng(); + let hashed_post_state_factory = HashedPostStateFactory::raw_transfer(); + let trie_updates_factory = TrieUpdatesFactory::raw_transfer(); + (0..len) + .map(|_| ExecutedBlockWithTrieUpdates { + block: ExecutedBlock { + recovered_block: Arc::default(), + execution_output: Arc::default(), + hashed_state: Arc::new(hashed_post_state_factory.generate(&mut rng)), + }, + trie: Arc::new(trie_updates_factory.generate(&mut rng)), + }) + .collect() +} + +fn generate_uniswap_blocks(len: usize) -> Vec> { + let mut rng = rand::thread_rng(); + let hashed_post_state_factory = HashedPostStateFactory::uniswap(); + let trie_updates_factory = TrieUpdatesFactory::uniswap(); + (0..len) + .map(|_| ExecutedBlockWithTrieUpdates { + block: ExecutedBlock { + recovered_block: Arc::default(), + execution_output: Arc::default(), + hashed_state: Arc::new(hashed_post_state_factory.generate(&mut rng)), + }, + trie: Arc::new(trie_updates_factory.generate(&mut rng)), + }) + .collect() +} + +#[inline] +fn run_trie_state(executed_blocks: Vec>) { + let provider = + MemoryOverlayStateProviderRef::::new(Box::new(NoopProvider::default()), executed_blocks); + + // Because `trie_state` is a private function, instead of calling it, we have to do this: + provider.storage_root(Address::ZERO, HashedStorage::default()).unwrap(); +} + +fn criterion_benchmark(c: &mut Criterion) { + c.bench_function("ERC20", |b| { + b.iter_custom(|iters| { + let mut total = Duration::ZERO; + for _ in 0..iters { + let blocks = generate_erc20_blocks(11); + let start = Instant::now(); + run_trie_state::(black_box(blocks)); + total += start.elapsed(); + } + total + }) + }); + c.bench_function("Raw Transfer", |b| { + b.iter_custom(|iters| { + let mut total = Duration::ZERO; + for _ in 0..iters { + let blocks = generate_raw_transfer_blocks(10); + let start = Instant::now(); + run_trie_state::(black_box(blocks)); + total += start.elapsed(); + } + total + }) + }); + c.bench_function("Uniswap", |b| { + b.iter_custom(|iters| { + let mut total = Duration::ZERO; + for _ in 0..iters { + let blocks = generate_uniswap_blocks(3); + let start = Instant::now(); + run_trie_state::(black_box(blocks)); + total += start.elapsed(); + } + total + }) + }); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches);