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