1use clap::Parser;
  2use miden_crypto::{
  3    hash::rpo::{Rpo256, RpoDigest},
  4    merkle::{MerkleError, Smt},
  5    Felt, Word, ONE,
  6};
  7use rand_utils::rand_value;
  8use std::time::Instant;
  9
 10#[derive(Parser, Debug)]
 11#[clap(name = "Benchmark", about = "SMT benchmark", version, rename_all = "kebab-case")]
 12pub struct BenchmarkCmd {
 13    /// Size of the tree
 14    #[clap(short = 's', long = "size")]
 15    size: u64,
 16}
 17
 18fn main() {
 19    benchmark_smt();
 20}
 21
 22/// Run a benchmark for [`Smt`].
 23pub fn benchmark_smt() {
 24    let args = BenchmarkCmd::parse();
 25    let tree_size = args.size;
 26
 27    // prepare the `leaves` vector for tree creation
 28    let mut entries = Vec::new();
 29    for i in 0..tree_size {
 30        let key = rand_value::<RpoDigest>();
 31        let value = [ONE, ONE, ONE, Felt::new(i)];
 32        entries.push((key, value));
 33    }
 34
 35    let mut tree = construction(entries, tree_size).unwrap();
 36    insertion(&mut tree, tree_size).unwrap();
 37    proof_generation(&mut tree, tree_size).unwrap();
 38}
 39
 40/// Runs the construction benchmark for [`Smt`], returning the constructed tree.
 41pub fn construction(entries: Vec<(RpoDigest, Word)>, size: u64) -> Result<Smt, MerkleError> {
 42    println!("Running a construction benchmark:");
 43    let now = Instant::now();
 44    let tree = Smt::with_entries(entries)?;
 45    let elapsed = now.elapsed();
 46    println!(
 47        "Constructed a SMT with {} key-value pairs in {:.3} seconds",
 48        size,
 49        elapsed.as_secs_f32(),
 50    );
 51
 52    println!("Number of leaf nodes: {}\n", tree.leaves().count());
 53
 54    Ok(tree)
 55}
 56
 57/// Runs the insertion benchmark for the [`Smt`].
 58pub fn insertion(tree: &mut Smt, size: u64) -> Result<(), MerkleError> {
 59    println!("Running an insertion benchmark:");
 60
 61    let mut insertion_times = Vec::new();
 62
 63    for i in 0..20 {
 64        let test_key = Rpo256::hash(&rand_value::<u64>().to_be_bytes());
 65        let test_value = [ONE, ONE, ONE, Felt::new(size + i)];
 66
 67        let now = Instant::now();
 68        tree.insert(test_key, test_value);
 69        let elapsed = now.elapsed();
 70        insertion_times.push(elapsed.as_secs_f32());
 71    }
 72
 73    println!(
 74        "An average insertion time measured by 20 inserts into a SMT with {} key-value pairs is {:.3} milliseconds\n",
 75        size,
 76        // calculate the average by dividing by 20 and convert to milliseconds by multiplying by 
 77        // 1000. As a result, we can only multiply by 50
 78        insertion_times.iter().sum::<f32>() * 50f32,
 79    );
 80
 81    Ok(())
 82}
 83
 84/// Runs the proof generation benchmark for the [`Smt`].
 85pub fn proof_generation(tree: &mut Smt, size: u64) -> Result<(), MerkleError> {
 86    println!("Running a proof generation benchmark:");
 87
 88    let mut insertion_times = Vec::new();
 89
 90    for i in 0..20 {
 91        let test_key = Rpo256::hash(&rand_value::<u64>().to_be_bytes());
 92        let test_value = [ONE, ONE, ONE, Felt::new(size + i)];
 93        tree.insert(test_key, test_value);
 94
 95        let now = Instant::now();
 96        let _proof = tree.open(&test_key);
 97        let elapsed = now.elapsed();
 98        insertion_times.push(elapsed.as_secs_f32());
 99    }
100
101    println!(
102        "An average proving time measured by 20 value proofs in a SMT with {} key-value pairs in {:.3} microseconds",
103        size,
104        // calculate the average by dividing by 20 and convert to microseconds by multiplying by
105        // 1000000. As a result, we can only multiply by 50000
106        insertion_times.iter().sum::<f32>() * 50000f32,
107    );
108
109    Ok(())
110}