feat(mpt): TrieNode benchmarks (#351)

Adds criterion benchmarks for `TrieNode` operations in the `kona-mpt`
crate.

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

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::<Vec<_>>();
+
+        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::<Vec<_>>();
+
+        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::<Vec<_>>();
+        let mut trie = TrieNode::Empty;
+
+        let rng = &mut rand::thread_rng();
+        let keys_to_delete = keys.choose_multiple(rng, 16).cloned().collect::<Vec<_>>();
+
+        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::<Vec<_>>();
+        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::<Vec<_>>();
+
+        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::<Vec<_>>();
+        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::<Vec<_>>();
+
+        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::<Vec<_>>();
+        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::<Vec<_>>();
+
+        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::<Vec<_>>();
+        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::<Vec<_>>();
+        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);

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