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