1use std::time::Instant;
  2
  3use clap::Parser;
  4use miden_crypto::{
  5    hash::rpo::{Rpo256, RpoDigest},
  6    merkle::{MerkleError, Smt},
  7    Felt, Word, ONE,
  8};
  9use rand_utils::rand_value;
 10
 11#[derive(Parser, Debug)]
 12#[clap(name = "Benchmark", about = "SMT benchmark", version, rename_all = "kebab-case")]
 13pub struct BenchmarkCmd {
 14    /// Size of the tree
 15    #[clap(short = 's', long = "size")]
 16    size: u64,
 17}
 18
 19fn main() {
 20    benchmark_smt();
 21}
 22
 23/// Run a benchmark for [`Smt`].
 24pub fn benchmark_smt() {
 25    let args = BenchmarkCmd::parse();
 26    let tree_size = args.size;
 27
 28    // prepare the `leaves` vector for tree creation
 29    let mut entries = Vec::new();
 30    for i in 0..tree_size {
 31        let key = rand_value::<RpoDigest>();
 32        let value = [ONE, ONE, ONE, Felt::new(i)];
 33        entries.push((key, value));
 34    }
 35
 36    let mut tree = construction(entries, tree_size).unwrap();
 37    insertion(&mut tree, tree_size).unwrap();
 38    batched_insertion(&mut tree, tree_size).unwrap();
 39    proof_generation(&mut tree, tree_size).unwrap();
 40}
 41
 42/// Runs the construction benchmark for [`Smt`], returning the constructed tree.
 43pub fn construction(entries: Vec<(RpoDigest, Word)>, size: u64) -> Result<Smt, MerkleError> {
 44    println!("Running a construction benchmark:");
 45    let now = Instant::now();
 46    let tree = Smt::with_entries(entries)?;
 47    let elapsed = now.elapsed();
 48    println!(
 49        "Constructed a SMT with {} key-value pairs in {:.3} seconds",
 50        size,
 51        elapsed.as_secs_f32(),
 52    );
 53
 54    println!("Number of leaf nodes: {}\n", tree.leaves().count());
 55
 56    Ok(tree)
 57}
 58
 59/// Runs the insertion benchmark for the [`Smt`].
 60pub fn insertion(tree: &mut Smt, size: u64) -> Result<(), MerkleError> {
 61    println!("Running an insertion benchmark:");
 62
 63    let mut insertion_times = Vec::new();
 64
 65    for i in 0..20 {
 66        let test_key = Rpo256::hash(&rand_value::<u64>().to_be_bytes());
 67        let test_value = [ONE, ONE, ONE, Felt::new(size + i)];
 68
 69        let now = Instant::now();
 70        tree.insert(test_key, test_value);
 71        let elapsed = now.elapsed();
 72        insertion_times.push(elapsed.as_secs_f32());
 73    }
 74
 75    println!(
 76        "An average insertion time measured by 20 inserts into a SMT with {} key-value pairs is {:.3} milliseconds\n",
 77        size,
 78        // calculate the average by dividing by 20 and convert to milliseconds by multiplying by
 79        // 1000. As a result, we can only multiply by 50
 80        insertion_times.iter().sum::<f32>() * 50f32,
 81    );
 82
 83    Ok(())
 84}
 85
 86pub fn batched_insertion(tree: &mut Smt, size: u64) -> Result<(), MerkleError> {
 87    println!("Running a batched insertion benchmark:");
 88
 89    let new_pairs: Vec<(RpoDigest, Word)> = (0..1000)
 90        .map(|i| {
 91            let key = Rpo256::hash(&rand_value::<u64>().to_be_bytes());
 92            let value = [ONE, ONE, ONE, Felt::new(size + i)];
 93            (key, value)
 94        })
 95        .collect();
 96
 97    let now = Instant::now();
 98    let mutations = tree.compute_mutations(new_pairs);
 99    let compute_elapsed = now.elapsed();
100
101    let now = Instant::now();
102    tree.apply_mutations(mutations).unwrap();
103    let apply_elapsed = now.elapsed();
104
105    println!(
106        "An average batch computation time measured by a 1k-batch into an SMT with {} key-value pairs over {:.3} milliseconds is {:.3} milliseconds",
107        size,
108        compute_elapsed.as_secs_f32() * 1000f32,
109        // Dividing by the number of iterations, 1000, and then multiplying by 1000 to get
110        // milliseconds, cancels out.
111        compute_elapsed.as_secs_f32(),
112    );
113
114    println!(
115        "An average batch application time measured by a 1k-batch into an SMT with {} key-value pairs over {:.3} milliseconds is {:.3} milliseconds",
116        size,
117        apply_elapsed.as_secs_f32() * 1000f32,
118        // Dividing by the number of iterations, 1000, and then multiplying by 1000 to get
119        // milliseconds, cancels out.
120        apply_elapsed.as_secs_f32(),
121    );
122
123    println!(
124        "An average batch insertion time measured by a 1k-batch into an SMT with {} key-value pairs totals to {:.3} milliseconds",
125        size,
126        (compute_elapsed + apply_elapsed).as_secs_f32() * 1000f32,
127    );
128
129    println!();
130
131    Ok(())
132}
133
134/// Runs the proof generation benchmark for the [`Smt`].
135pub fn proof_generation(tree: &mut Smt, size: u64) -> Result<(), MerkleError> {
136    println!("Running a proof generation benchmark:");
137
138    let mut insertion_times = Vec::new();
139
140    for i in 0..20 {
141        let test_key = Rpo256::hash(&rand_value::<u64>().to_be_bytes());
142        let test_value = [ONE, ONE, ONE, Felt::new(size + i)];
143        tree.insert(test_key, test_value);
144
145        let now = Instant::now();
146        let _proof = tree.open(&test_key);
147        let elapsed = now.elapsed();
148        insertion_times.push(elapsed.as_secs_f32());
149    }
150
151    println!(
152        "An average proving time measured by 20 value proofs in a SMT with {} key-value pairs in {:.3} microseconds",
153        size,
154        // calculate the average by dividing by 20 and convert to microseconds by multiplying by
155        // 1000000. As a result, we can only multiply by 50000
156        insertion_times.iter().sum::<f32>() * 50000f32,
157    );
158
159    Ok(())
160}