From ae191003a592a41897305ef045258b9068636ef0 Mon Sep 17 00:00:00 2001
From: Qyriad <qyriad@qyriad.me>
Date: Thu, 14 Nov 2024 14:45:26 -0700
Subject: [PATCH] merkle: add a benchmark for constructing 256-balanced trees

This is intended for comparison with the benchmarks from the previous
commit. This benchmark represents the theoretical perfect-efficiency
performance we could possibly (but impractically) get for computing
depth-8 sparse Merkle subtrees.
---
 Cargo.toml        |  4 +++
 benches/merkle.rs | 66 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+)
 create mode 100644 benches/merkle.rs

diff --git a/Cargo.toml b/Cargo.toml
index ec59d42..74df3ab 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -31,6 +31,10 @@ harness = false
 name = "smt-subtree"
 harness = false
 
+[[bench]]
+name = "merkle"
+harness = false
+
 [[bench]]
 name = "store"
 harness = false
diff --git a/benches/merkle.rs b/benches/merkle.rs
new file mode 100644
index 0000000..7d6bb2c
--- /dev/null
+++ b/benches/merkle.rs
@@ -0,0 +1,66 @@
+//! Benchmark for building a [`miden_crypto::merkle::MerkleTree`]. This is intended to be compared
+//! with the results from `benches/smt-subtree.rs`, as building a fully balanced Merkle tree with
+//! 256 leaves should indicate the *absolute best* performance we could *possibly* get for building
+//! a depth-8 sparse Merkle subtree, though practically speaking building a fully balanced Merkle
+//! tree will perform better than the sparse version. At the time of this writing (2024/11/24), this
+//! benchmark is about four times more efficient than the equivalent benchmark in
+//! `benches/smt-subtree.rs`.
+use std::{hint, mem, time::Duration};
+
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
+use miden_crypto::{merkle::MerkleTree, Felt, Word, ONE};
+use rand_utils::prng_array;
+
+fn balanced_merkle_even(c: &mut Criterion) {
+    c.bench_function("balanced-merkle-even", |b| {
+        b.iter_batched(
+            || {
+                let entries: Vec<Word> =
+                    (0..256).map(|i| [Felt::new(i), ONE, ONE, Felt::new(i)]).collect();
+                assert_eq!(entries.len(), 256);
+                entries
+            },
+            |leaves| {
+                let tree = MerkleTree::new(hint::black_box(leaves)).unwrap();
+                assert_eq!(tree.depth(), 8);
+            },
+            BatchSize::SmallInput,
+        );
+    });
+}
+
+fn balanced_merkle_rand(c: &mut Criterion) {
+    let mut seed = [0u8; 32];
+    c.bench_function("balanced-merkle-rand", |b| {
+        b.iter_batched(
+            || {
+                let entries: Vec<Word> = (0..256).map(|_| generate_word(&mut seed)).collect();
+                assert_eq!(entries.len(), 256);
+                entries
+            },
+            |leaves| {
+                let tree = MerkleTree::new(hint::black_box(leaves)).unwrap();
+                assert_eq!(tree.depth(), 8);
+            },
+            BatchSize::SmallInput,
+        );
+    });
+}
+
+criterion_group! {
+    name = smt_subtree_group;
+    config = Criterion::default()
+        .measurement_time(Duration::from_secs(20))
+        .configure_from_args();
+    targets = balanced_merkle_even, balanced_merkle_rand
+}
+criterion_main!(smt_subtree_group);
+
+// HELPER FUNCTIONS
+// --------------------------------------------------------------------------------------------
+
+fn generate_word(seed: &mut [u8; 32]) -> Word {
+    mem::swap(seed, &mut prng_array(*seed));
+    let nums: [u64; 4] = prng_array(*seed);
+    [Felt::new(nums[0]), Felt::new(nums[1]), Felt::new(nums[2]), Felt::new(nums[3])]
+}