test: update and fix Merkle benchmarks (#5964)

Description
---
Updates and fixes Merkle benchmarks.

Closes #5962.

Motivation and Context
---
Recent work in #5954 removes mutable Merkle mountain range (MMR) code,
which became unused with the addition of sparse Merkle trees (SMTs).
Some existing benchmarks were removed or moved to accommodate this
change. During review, some minor issues were identified:
- An existing MMR and SMT benchmark reused tree structures between
benchmark iterations, which could provide incorrect timing data.
- New SMT benchmarks used bespoke timing code that didn't take full
advantage of Criterion's functionality.

This PR fixes these issues, albeit at the expense of more benchmark
setup overhead. This shouldn't be particularly problematic, as
benchmarks are only run manually as needed.

How Has This Been Tested?
---
The benchmarks run and appear to give reasonable data.

What process can a PR reviewer use to test or verify this change?
---
Check that the updated and new benchmarks exercise the desired
functionality. Check that operations use fresh tree structures for each
test iteration as needed, to ensure correct timing data.

base_layer/mmr/benches/mmr.rs

@@ -37,14 +37,16 @@ fn get_hashes(n: usize) -> Vec<Vec<u8>> {
 }
 
 fn build_mmr(c: &mut Criterion) {
-    let sizes = [100, 10_000];
+    let sizes = [100, 1_000, 10_000, 100_000];
     for size in sizes {
         c.bench_function(&format!("MMR: {size} hashes"), move |b| {
             let hashes = get_hashes(size);
-            let mut mmr = TestMmr::new(Vec::default());
             b.iter_batched(
-                || hashes.clone(),
-                |hashes| {
+                || {
+                    // Set up a fresh tree for this iteration
+                    (TestMmr::new(Vec::default()), hashes.clone())
+                },
+                |(mut mmr, hashes)| {
                     hashes.into_iter().for_each(|hash| {
                         mmr.push(hash).unwrap();
                     });

base_layer/mmr/benches/smt.rs

@@ -6,89 +6,158 @@ use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
 use digest::consts::U32;
 use tari_mmr::sparse_merkle_tree::{NodeKey, SparseMerkleTree, ValueHash};
 
+type TestSmt = SparseMerkleTree<Blake2b<U32>>;
+
+// The number of keys to use for full trees
+const SIZES: [usize; 4] = [100, 1_000, 10_000, 100_000];
+
+// Helper to generate a single random key
 fn random_key() -> NodeKey {
     let key = rand::random::<[u8; 32]>();
     NodeKey::from(key)
 }
 
+// Helper to generate a set of random keys
 fn get_keys(n: usize) -> Vec<NodeKey> {
     (0..n).map(|_| random_key()).collect()
 }
 
-fn create_smt() -> SparseMerkleTree<Blake2b<U32>> {
-    SparseMerkleTree::<Blake2b<U32>>::new()
+// Helper to upsert keys
+fn upsert_keys(smt: &mut TestSmt, keys: Vec<NodeKey>) {
+    keys.into_iter().for_each(|key| {
+        smt.upsert(key, ValueHash::default()).unwrap();
+    });
+}
+
+// Helper to delete keys
+fn delete_keys(smt: &mut TestSmt, keys: &[NodeKey]) {
+    keys.iter().for_each(|key| {
+        smt.delete(key).unwrap();
+    });
 }
 
-pub fn benchmark_smt_insert(c: &mut Criterion) {
-    let sizes = [100, 10_000];
-    for size in sizes {
+// Build an SMT by inserting keys
+pub fn build_smt(c: &mut Criterion) {
+    for size in SIZES {
         c.bench_function(&format!("SMT: Insert {size} keys"), move |b| {
             let keys = get_keys(size);
-            let mut smt = create_smt();
             b.iter_batched(
-                || keys.clone(),
-                |hashes| {
-                    hashes.into_iter().for_each(|key| {
-                        smt.upsert(key, ValueHash::default()).unwrap();
-                    });
+                || {
+                    // Set up a fresh tree for this iteration
+                    (TestSmt::new(), keys.clone())
+                },
+                |(mut smt, hashes)| {
+                    upsert_keys(&mut smt, hashes);
                 },
                 BatchSize::SmallInput,
             );
         });
     }
 }
 
-fn insert_into_smt(keys: &[NodeKey], tree: &mut SparseMerkleTree<Blake2b<U32>>) {
-    keys.iter().for_each(|key| {
-        tree.upsert(key.clone(), ValueHash::default()).unwrap();
-    });
+// Compute the root hash of a full tree
+pub fn full_root_hash(c: &mut Criterion) {
+    for size in SIZES {
+        c.bench_function(&format!("SMT: Full root hash on {size}-key tree"), move |b| {
+            // We can reuse the same tree between iterations
+            let keys = get_keys(size);
+            let mut smt = TestSmt::new();
+            upsert_keys(&mut smt, keys);
+
+            b.iter(|| {
+                smt.root();
+            });
+        });
+    }
 }
 
-fn delete_from_smt(keys: &[NodeKey], tree: &mut SparseMerkleTree<Blake2b<U32>>) {
-    keys.iter().for_each(|key| {
-        tree.delete(key).unwrap();
-    });
+// Delete half of the keys from a full tree
+pub fn delete_half_keys(c: &mut Criterion) {
+    for size in SIZES {
+        c.bench_function(&format!("SMT: Delete half of keys on {size}-key tree"), move |b| {
+            let keys = get_keys(size);
+            b.iter_batched(
+                || {
+                    // Build a a fresh tree for this iteration
+                    let mut smt = TestSmt::new();
+                    upsert_keys(&mut smt, keys.clone());
+
+                    (smt, keys.clone())
+                },
+                |(mut smt, keys)| {
+                    delete_keys(&mut smt, &keys[..size / 2]);
+                },
+                BatchSize::SmallInput,
+            );
+        });
+    }
 }
 
-fn time_function(header: &str, f: impl FnOnce()) -> std::time::Duration {
-    println!("Starting: {header}");
-    let now = std::time::Instant::now();
-    f();
-    let t = now.elapsed();
-    println!("Finished: {header} - {t:?}");
-    t
+// Compute the root hash of a half-empty tree
+pub fn half_root_hash(c: &mut Criterion) {
+    for size in SIZES {
+        c.bench_function(&format!("SMT: Half-empty root hash on {size}-key tree"), move |b| {
+            // We can reuse the same tree between iterations
+            let keys = get_keys(size);
+            let mut smt = TestSmt::new();
+            upsert_keys(&mut smt, keys.clone());
+            delete_keys(&mut smt, &keys[..size / 2]);
+
+            b.iter(|| {
+                smt.root();
+            });
+        });
+    }
 }
 
-pub fn root_hash(_c: &mut Criterion) {
-    let size = 1_000_000;
-    let half_size = size / 2;
-    let keys = get_keys(size);
-    let mut tree = create_smt();
-    time_function(&format!("SMT: Inserting {size} keys"), || {
-        insert_into_smt(&keys, &mut tree);
-    });
-    time_function("SMT: Calculating root hash", || {
-        let size = tree.size();
-        let hash = tree.hash();
-        println!("Tree size: {size}. Root hash: {hash:x}");
-    });
-    time_function(&format!("SMT: Deleting {half_size} keys"), || {
-        delete_from_smt(&keys[0..half_size], &mut tree);
-    });
-    time_function("SMT: Calculating root hash", || {
-        let size = tree.size();
-        let hash = tree.hash();
-        println!("Tree size: {size}. Root hash: {hash:x}");
-    });
-    time_function(&format!("SMT: Deleting another {half_size} keys"), || {
-        delete_from_smt(&keys[half_size..], &mut tree);
-    });
-    time_function("SMT: Calculating root hash", || {
-        let size = tree.size();
-        let hash = tree.hash();
-        println!("Tree size: {size}. Root hash: {hash:x}");
-    });
+// Delete remaining half of the keys from a half-empty tree
+pub fn delete_remaining_keys(c: &mut Criterion) {
+    for size in SIZES {
+        c.bench_function(&format!("SMT: Delete half of keys on {size}-key tree"), move |b| {
+            let keys = get_keys(size);
+            b.iter_batched(
+                || {
+                    // Build a a fresh tree for this iteration
+                    let mut smt = TestSmt::new();
+                    upsert_keys(&mut smt, keys.clone());
+                    delete_keys(&mut smt, &keys[..size / 2]);
+
+                    (smt, keys.clone())
+                },
+                |(mut smt, keys)| {
+                    delete_keys(&mut smt, &keys[size / 2..]);
+                },
+                BatchSize::SmallInput,
+            );
+        });
+    }
+}
+
+// Compute the root hash of an empty tree
+pub fn empty_root_hash(c: &mut Criterion) {
+    for size in SIZES {
+        c.bench_function(&format!("SMT: Half-empty root hash on {size}-key tree"), move |b| {
+            // We can reuse the same tree between iterations
+            let keys = get_keys(size);
+            let mut smt = TestSmt::new();
+            upsert_keys(&mut smt, keys.clone());
+            delete_keys(&mut smt, &keys[..size / 2]);
+            delete_keys(&mut smt, &keys[size / 2..]);
+
+            b.iter(|| {
+                smt.root();
+            });
+        });
+    }
 }
 
-criterion_group!(smt, benchmark_smt_insert, root_hash);
+criterion_group!(
+    smt,
+    build_smt,
+    full_root_hash,
+    delete_half_keys,
+    half_root_hash,
+    delete_remaining_keys,
+    empty_root_hash
+);
 criterion_main!(smt);