multiproof bench code

multiproof_allocs.rs

#![allow(missing_docs, unreachable_pub)]
use alloy_primitives::{map::B256Set, B256};
use criterion::{black_box, criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
use reth_trie_common::MultiProofTargets;

fn b256_from_u64(value: u64) -> B256 {
    let mut bytes = [0u8; 32];
    bytes[24..].copy_from_slice(&value.to_be_bytes());
    B256::from(bytes)
}

fn make_targets(accounts: usize, slots_per_account: usize, start: u64) -> MultiProofTargets {
    let mut targets = MultiProofTargets::with_capacity(accounts);
    for i in 0..accounts {
        let address = b256_from_u64(start + i as u64);
        let mut slots = B256Set::default();
        for j in 0..slots_per_account {
            let slot_key = (start as u128) << 32 | (i as u128) << 16 | j as u128;
            slots.insert(b256_from_u64(slot_key as u64));
        }
        targets.insert(address, slots);
    }
    targets
}

fn bench_multiproof_targets_extend(c: &mut Criterion) {
    let mut group = c.benchmark_group("MultiProofTargets::extend");
    for (accounts, slots_per_account) in [(256, 4), (1024, 2), (2048, 1)] {
        let id = format!("accounts_{accounts}/slots_{slots_per_account}");
        group.bench_function(BenchmarkId::new("extend_owned", &id), |b| {
            b.iter_batched(
                || {
                    let acc = make_targets(accounts, slots_per_account, 0);
                    let other = make_targets(accounts, slots_per_account, 10_000);
                    (acc, other)
                },
                |(mut acc, other)| {
                    acc.extend(other);
                    black_box(acc);
                },
                BatchSize::SmallInput,
            );
        });
    }
}

criterion_group!(multiproof_allocs, bench_multiproof_targets_extend);
criterion_main!(multiproof_allocs);