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    /// Run the construction benchmark
 24    #[clap(short = 'c', long = "construction")]
 25    construction: bool,
 26
 27    /// Run the insertion benchmark
 28    #[clap(short = 'i', long = "insertion")]
 29    insertion: bool,
 30
 31    /// Run the proof generation benchmark
 32    #[clap(short = 'p', long = "proof-generation")]
 33    proof_generation: bool,
 34}
 35
 36fn main() {
 37    benchmark_tsmt();
 38}
 39
 40/// Run a benchmark for the Tiered SMT.
 41pub fn benchmark_tsmt() {
 42    let args = BenchmarkCmd::parse();
 43    let tree_size = args.size;
 44
 45    // prepare the `leaves` vector for tree creation
 46    let mut leaves = Vec::new();
 47    for i in 0..tree_size {
 48        let key = rand_value::<RpoDigest>();
 49        let value = [ONE, ONE, ONE, Felt::new(i)];
 50        leaves.push((key, value));
 51    }
 52
 53    let mut tree: Option<TieredSmt> = None;
 54
 55    // if the `-c` argument was specified
 56    if args.construction {
 57        tree = Some(construction(leaves.clone(), tree_size).unwrap());
 58    }
 59
 60    // if the `-i` argument was specified
 61    if args.insertion {
 62        if let Some(inner_tree) = tree {
 63            tree = Some(insertion(inner_tree, tree_size).unwrap());
 64        } else {
 65            let inner_tree = TieredSmt::with_leaves(leaves.clone()).unwrap();
 66            tree = Some(insertion(inner_tree, tree_size).unwrap());
 67        }
 68    }
 69
 70    // if the `-p` argument was specified
 71    if args.proof_generation {
 72        if let Some(inner_tree) = tree {
 73            proof_generation(inner_tree, tree_size).unwrap();
 74        } else {
 75            let inner_tree = TieredSmt::with_leaves(leaves).unwrap();
 76            proof_generation(inner_tree, tree_size).unwrap();
 77        }
 78    }
 79}
 80
 81/// Run the construction benchmark for the Tiered SMT.
 82pub fn construction(leaves: Vec<(RpoDigest, Word)>, size: u64) -> Result<TieredSmt, MerkleError> {
 83    println!("Running a construction benchmark:");
 84    let now = Instant::now();
 85    let tree = TieredSmt::with_leaves(leaves)?;
 86    let elapsed = now.elapsed();
 87    println!(
 88        "Constructed a TSMT with {} key-value pairs in {:.3} seconds",
 89        size,
 90        elapsed.as_secs_f32(),
 91    );
 92
 93    // Count how many nodes end up at each tier
 94    let mut nodes_num_16_32_48 = (0, 0, 0);
 95
 96    tree.upper_leaf_nodes().for_each(|(index, _)| match index.depth() {
 97        16 => nodes_num_16_32_48.0 += 1,
 98        32 => nodes_num_16_32_48.1 += 1,
 99        48 => nodes_num_16_32_48.2 += 1,
100        _ => unreachable!(),
101    });
102
103    println!("Number of nodes on depth 16: {}", nodes_num_16_32_48.0);
104    println!("Number of nodes on depth 32: {}", nodes_num_16_32_48.1);
105    println!("Number of nodes on depth 48: {}", nodes_num_16_32_48.2);
106    println!("Number of nodes on depth 64: {}\n", tree.bottom_leaves().count());
107
108    Ok(tree)
109}
110
111/// Run the insertion benchmark for the Tiered SMT.
112pub fn insertion(mut tree: TieredSmt, size: u64) -> Result<TieredSmt, MerkleError> {
113    println!("Running an insertion benchmark:");
114
115    let mut insertion_times = Vec::new();
116
117    for i in 0..20 {
118        let test_key = Rpo256::hash(&rand_value::<u64>().to_be_bytes());
119        let test_value = [ONE, ONE, ONE, Felt::new(size + i)];
120
121        let now = Instant::now();
122        tree.insert(test_key, test_value);
123        let elapsed = now.elapsed();
124        insertion_times.push(elapsed.as_secs_f32());
125    }
126
127    println!(
128        "An average insertion time measured by 20 inserts into a TSMT with {} key-value pairs is {:.3} milliseconds\n",
129        size,
130        // calculate the average by dividing by 20 and convert to milliseconds by multiplying by 
131        // 1000. As a result, we can only multiply by 50
132        insertion_times.iter().sum::<f32>() * 50f32,
133    );
134
135    Ok(tree)
136}
137
138/// Run the proof generation benchmark for the Tiered SMT.
139pub fn proof_generation(mut tree: TieredSmt, size: u64) -> Result<(), MerkleError> {
140    println!("Running a proof generation benchmark:");
141
142    let mut insertion_times = Vec::new();
143
144    for i in 0..20 {
145        let test_key = Rpo256::hash(&rand_value::<u64>().to_be_bytes());
146        let test_value = [ONE, ONE, ONE, Felt::new(size + i)];
147        tree.insert(test_key, test_value);
148
149        let now = Instant::now();
150        let _proof = tree.prove(test_key);
151        let elapsed = now.elapsed();
152        insertion_times.push(elapsed.as_secs_f32());
153    }
154
155    println!(
156        "An average proving time measured by 20 value proofs in a TSMT with {} key-value pairs in {:.3} microseconds",
157        size,
158        // calculate the average by dividing by 20 and convert to microseconds by multiplying by
159        // 1000000. As a result, we can only multiply by 50000
160        insertion_times.iter().sum::<f32>() * 50000f32,
161    );
162
163    Ok(())
164}