From d14e7f51a814e072ffd3373040a9d8aa6ed677ec Mon Sep 17 00:00:00 2001 From: clabby Date: Mon, 8 Jul 2024 13:02:23 -0400 Subject: [PATCH] feat(mpt): `TrieNode` benchmarks (#351) Adds criterion benchmarks for `TrieNode` operations in the `kona-mpt` crate. --- Cargo.lock | 2 + crates/mpt/Cargo.toml | 7 ++ crates/mpt/benches/trie_node.rs | 149 ++++++++++++++++++++++++++++++++ crates/mpt/src/node.rs | 40 ++++++++- 4 files changed, 197 insertions(+), 1 deletion(-) create mode 100644 crates/mpt/benches/trie_node.rs diff --git a/Cargo.lock b/Cargo.lock index 3ba08e4e6..dad75ad72 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2238,7 +2238,9 @@ dependencies = [ "alloy-transport-http", "alloy-trie", "anyhow", + "criterion", "futures", + "pprof", "proptest", "rand", "reqwest", diff --git a/crates/mpt/Cargo.toml b/crates/mpt/Cargo.toml index c80acbc94..84da0285b 100644 --- a/crates/mpt/Cargo.toml +++ b/crates/mpt/Cargo.toml @@ -29,5 +29,12 @@ alloy-transport-http = { version = "0.1" } reqwest = "0.12.4" tracing-subscriber = "0.3.18" futures = { version = "0.3.30", default-features = false } + proptest = "1.4" rand = "0.8.5" +criterion = { version = "0.5.1", features = ["html_reports"] } +pprof = { version = "0.13.0", features = ["criterion", "flamegraph", "frame-pointer"] } + +[[bench]] +name = "trie_node" +harness = false diff --git a/crates/mpt/benches/trie_node.rs b/crates/mpt/benches/trie_node.rs new file mode 100644 index 000000000..c20f05ca8 --- /dev/null +++ b/crates/mpt/benches/trie_node.rs @@ -0,0 +1,149 @@ +//! Contains benchmarks for the [TrieNode]. + +use alloy_trie::Nibbles; +use criterion::{criterion_group, criterion_main, Criterion}; +use kona_mpt::{NoopTrieDBFetcher, NoopTrieDBHinter, TrieNode}; +use pprof::criterion::{Output, PProfProfiler}; +use rand::{rngs::StdRng, seq::SliceRandom, Rng, SeedableRng}; + +fn trie(c: &mut Criterion) { + let mut g = c.benchmark_group("execution"); + g.sample_size(10); + + // Use pseudo-randomness for reproducibility + let mut rng = StdRng::seed_from_u64(42); + + g.bench_function("Insertion - 4096 nodes", |b| { + let keys = + (0..2usize.pow(12)).map(|_| Nibbles::unpack(rng.gen::<[u8; 32]>())).collect::>(); + + b.iter(|| { + let mut trie = TrieNode::Empty; + for key in &keys { + trie.insert(key, key.to_vec().into(), &NoopTrieDBFetcher).unwrap(); + } + }); + }); + + g.bench_function("Insertion - 65,536 nodes", |b| { + let keys = + (0..2usize.pow(16)).map(|_| Nibbles::unpack(rng.gen::<[u8; 32]>())).collect::>(); + + b.iter(|| { + let mut trie = TrieNode::Empty; + for key in &keys { + trie.insert(key, key.to_vec().into(), &NoopTrieDBFetcher).unwrap(); + } + }); + }); + + g.bench_function("Delete 16 nodes - 4096 nodes", |b| { + let keys = + (0..2usize.pow(12)).map(|_| Nibbles::unpack(rng.gen::<[u8; 32]>())).collect::>(); + let mut trie = TrieNode::Empty; + + let rng = &mut rand::thread_rng(); + let keys_to_delete = keys.choose_multiple(rng, 16).cloned().collect::>(); + + for key in &keys { + trie.insert(key, key.to_vec().into(), &NoopTrieDBFetcher).unwrap(); + } + + b.iter(|| { + let trie = &mut trie.clone(); + for key in &keys_to_delete { + trie.delete(key, &NoopTrieDBFetcher, &NoopTrieDBHinter).unwrap(); + } + }); + }); + + g.bench_function("Delete 16 nodes - 65,536 nodes", |b| { + let keys = + (0..2usize.pow(16)).map(|_| Nibbles::unpack(rng.gen::<[u8; 32]>())).collect::>(); + let mut trie = TrieNode::Empty; + for key in &keys { + trie.insert(key, key.to_vec().into(), &NoopTrieDBFetcher).unwrap(); + } + + let rng = &mut rand::thread_rng(); + let keys_to_delete = keys.choose_multiple(rng, 16).cloned().collect::>(); + + b.iter(|| { + let trie = &mut trie.clone(); + for key in &keys_to_delete { + trie.delete(key, &NoopTrieDBFetcher, &NoopTrieDBHinter).unwrap(); + } + }); + }); + + g.bench_function("Open 1024 nodes - 4096 nodes", |b| { + let keys = + (0..2usize.pow(12)).map(|_| Nibbles::unpack(rng.gen::<[u8; 32]>())).collect::>(); + let mut trie = TrieNode::Empty; + for key in &keys { + trie.insert(key, key.to_vec().into(), &NoopTrieDBFetcher).unwrap(); + } + + let rng = &mut rand::thread_rng(); + let keys_to_retrieve = keys.choose_multiple(rng, 1024).cloned().collect::>(); + + b.iter(|| { + for key in &keys_to_retrieve { + trie.open(key, &NoopTrieDBFetcher).unwrap(); + } + }); + }); + + g.bench_function("Open 1024 nodes - 65,536 nodes", |b| { + let keys = + (0..2usize.pow(16)).map(|_| Nibbles::unpack(rng.gen::<[u8; 32]>())).collect::>(); + let mut trie = TrieNode::Empty; + for key in &keys { + trie.insert(key, key.to_vec().into(), &NoopTrieDBFetcher).unwrap(); + } + + let rng = &mut rand::thread_rng(); + let keys_to_retrieve = keys.choose_multiple(rng, 1024).cloned().collect::>(); + + b.iter(|| { + for key in &keys_to_retrieve { + trie.open(key, &NoopTrieDBFetcher).unwrap(); + } + }); + }); + + g.bench_function("Compute root, fully open trie - 4096 nodes", |b| { + let keys = + (0..2usize.pow(12)).map(|_| Nibbles::unpack(rng.gen::<[u8; 32]>())).collect::>(); + let mut trie = TrieNode::Empty; + for key in &keys { + trie.insert(key, key.to_vec().into(), &NoopTrieDBFetcher).unwrap(); + } + + b.iter(|| { + let trie = &mut trie.clone(); + trie.blind(); + }); + }); + + g.bench_function("Compute root, fully open trie - 65,536 nodes", |b| { + let keys = + (0..2usize.pow(16)).map(|_| Nibbles::unpack(rng.gen::<[u8; 32]>())).collect::>(); + let mut trie = TrieNode::Empty; + for key in &keys { + trie.insert(key, key.to_vec().into(), &NoopTrieDBFetcher).unwrap(); + } + + b.iter(|| { + let trie = &mut trie.clone(); + trie.blind(); + }); + }); +} + +criterion_group! { + name = trie_benches; + config = Criterion::default().with_profiler(PProfProfiler::new(100, Output::Flamegraph(None))); + targets = trie +} +criterion_main!(trie_benches); diff --git a/crates/mpt/src/node.rs b/crates/mpt/src/node.rs index 494d10728..63bcfd423 100644 --- a/crates/mpt/src/node.rs +++ b/crates/mpt/src/node.rs @@ -122,7 +122,7 @@ impl TrieNode { pub fn blind(&mut self) { if self.length() >= B256::ZERO.len() && !matches!(self, TrieNode::Blinded { .. }) { let mut rlp_buf = Vec::with_capacity(self.length()); - self.encode(&mut rlp_buf); + self.encode_in_place(&mut rlp_buf); *self = TrieNode::Blinded { commitment: keccak256(rlp_buf) } } } @@ -359,6 +359,44 @@ impl TrieNode { } } + /// Alternative function to the [Encodable::encode] implementation for this type, that blinds + /// children nodes throughout the encoding process. This function is useful in the case where + /// the trie node cache is no longer required (i.e., during [Self::blind]). + /// + /// ## Takes + /// - `self` - The root trie node + /// - `out` - The buffer to write the encoded trie node to + pub fn encode_in_place(&mut self, out: &mut dyn alloy_rlp::BufMut) { + let payload_length = self.payload_length(); + match self { + Self::Empty => out.put_u8(EMPTY_STRING_CODE), + Self::Blinded { commitment } => commitment.encode(out), + Self::Leaf { prefix, value } => { + // Encode the leaf node's header and key-value pair. + Header { list: true, payload_length }.encode(out); + prefix.encode_path_leaf(true).as_slice().encode(out); + value.encode(out); + } + Self::Extension { prefix, node } => { + // Encode the extension node's header, prefix, and pointer node. + Header { list: true, payload_length }.encode(out); + prefix.encode_path_leaf(false).as_slice().encode(out); + node.blind(); + node.encode_in_place(out); + } + Self::Branch { stack } => { + // In branch nodes, if an element is longer than 32 bytes in length, it is blinded. + // Assuming we have an open trie node, we must re-hash the elements + // that are longer than 32 bytes in length. + Header { list: true, payload_length }.encode(out); + stack.iter_mut().for_each(|node| { + node.blind(); + node.encode_in_place(out); + }); + } + } + } + /// If applicable, collapses `self` into a more compact form. /// /// ## Takes