1use std::time::Instant;
  2
  3use clap::Parser;
  4use miden_crypto::{
  5    hash::rpo::{Rpo256, RpoDigest},
  6    merkle::{MerkleError, Smt},
  7    Felt, Word, EMPTY_WORD, ONE,
  8};
  9use rand::{prelude::IteratorRandom, thread_rng, Rng};
 10use rand_utils::rand_value;
 11
 12#[derive(Parser, Debug)]
 13#[clap(name = "Benchmark", about = "SMT benchmark", version, rename_all = "kebab-case")]
 14pub struct BenchmarkCmd {
 15    /// Size of the tree
 16    #[clap(short = 's', long = "size")]
 17    size: usize,
 18}
 19
 20fn main() {
 21    benchmark_smt();
 22}
 23
 24/// Run a benchmark for [`Smt`].
 25pub fn benchmark_smt() {
 26    let args = BenchmarkCmd::parse();
 27    let tree_size = args.size;
 28
 29    // prepare the `leaves` vector for tree creation
 30    let mut entries = Vec::new();
 31    for i in 0..tree_size {
 32        let key = rand_value::<RpoDigest>();
 33        let value = [ONE, ONE, ONE, Felt::new(i as u64)];
 34        entries.push((key, value));
 35    }
 36
 37    let mut tree = construction(entries.clone(), tree_size).unwrap();
 38    insertion(&mut tree).unwrap();
 39    batched_insertion(&mut tree).unwrap();
 40    batched_update(&mut tree, entries).unwrap();
 41    proof_generation(&mut tree).unwrap();
 42}
 43
 44/// Runs the construction benchmark for [`Smt`], returning the constructed tree.
 45pub fn construction(entries: Vec<(RpoDigest, Word)>, size: usize) -> Result<Smt, MerkleError> {
 46    println!("Running a construction benchmark:");
 47    let now = Instant::now();
 48    let tree = Smt::with_entries(entries)?;
 49    let elapsed = now.elapsed().as_secs_f32();
 50
 51    println!("Constructed a SMT with {size} key-value pairs in {elapsed:.1} seconds");
 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) -> Result<(), MerkleError> {
 59    const NUM_INSERTIONS: usize = 1_000;
 60
 61    println!("Running an insertion benchmark:");
 62
 63    let size = tree.num_leaves();
 64
 65    let mut insertion_times = Vec::new();
 66
 67    for i in 0..NUM_INSERTIONS {
 68        let test_key = Rpo256::hash(&rand_value::<u64>().to_be_bytes());
 69        let test_value = [ONE, ONE, ONE, Felt::new((size + i) as u64)];
 70
 71        let now = Instant::now();
 72        tree.insert(test_key, test_value);
 73        let elapsed = now.elapsed();
 74        insertion_times.push(elapsed.as_micros());
 75    }
 76
 77    println!(
 78        "An average insertion time measured by {NUM_INSERTIONS} inserts into an SMT with {size} leaves is {:.0} μs\n",
 79        // calculate the average
 80        insertion_times.iter().sum::<u128>() as f64 / (NUM_INSERTIONS as f64),
 81    );
 82
 83    Ok(())
 84}
 85
 86pub fn batched_insertion(tree: &mut Smt) -> Result<(), MerkleError> {
 87    const NUM_INSERTIONS: usize = 1_000;
 88
 89    println!("Running a batched insertion benchmark:");
 90
 91    let size = tree.num_leaves();
 92
 93    let new_pairs: Vec<(RpoDigest, Word)> = (0..NUM_INSERTIONS)
 94        .map(|i| {
 95            let key = Rpo256::hash(&rand_value::<u64>().to_be_bytes());
 96            let value = [ONE, ONE, ONE, Felt::new((size + i) as u64)];
 97            (key, value)
 98        })
 99        .collect();
100
101    let now = Instant::now();
102    let mutations = tree.compute_mutations(new_pairs);
103    let compute_elapsed = now.elapsed().as_secs_f64() * 1000_f64; // time in ms
104
105    let now = Instant::now();
106    tree.apply_mutations(mutations)?;
107    let apply_elapsed = now.elapsed().as_secs_f64() * 1000_f64; // time in ms
108
109    println!(
110        "An average insert-batch computation time measured by a {NUM_INSERTIONS}-batch into an SMT with {size} leaves over {:.1} ms is {:.0} μs",
111        compute_elapsed,
112        compute_elapsed * 1000_f64 / NUM_INSERTIONS as f64, // time in μs
113    );
114
115    println!(
116        "An average insert-batch application time measured by a {NUM_INSERTIONS}-batch into an SMT with {size} leaves over {:.1} ms is {:.0} μs",
117        apply_elapsed,
118        apply_elapsed * 1000_f64 / NUM_INSERTIONS as f64, // time in μs
119    );
120
121    println!(
122        "An average batch insertion time measured by a 1k-batch into an SMT with {size} leaves totals to {:.1} ms",
123        (compute_elapsed + apply_elapsed),
124    );
125
126    println!();
127
128    Ok(())
129}
130
131pub fn batched_update(tree: &mut Smt, entries: Vec<(RpoDigest, Word)>) -> Result<(), MerkleError> {
132    const NUM_UPDATES: usize = 1_000;
133    const REMOVAL_PROBABILITY: f64 = 0.2;
134
135    println!("Running a batched update benchmark:");
136
137    let size = tree.num_leaves();
138    let mut rng = thread_rng();
139
140    let new_pairs =
141        entries
142            .into_iter()
143            .choose_multiple(&mut rng, NUM_UPDATES)
144            .into_iter()
145            .map(|(key, _)| {
146                let value = if rng.gen_bool(REMOVAL_PROBABILITY) {
147                    EMPTY_WORD
148                } else {
149                    [ONE, ONE, ONE, Felt::new(rng.gen())]
150                };
151
152                (key, value)
153            });
154
155    assert_eq!(new_pairs.len(), NUM_UPDATES);
156
157    let now = Instant::now();
158    let mutations = tree.compute_mutations(new_pairs);
159    let compute_elapsed = now.elapsed().as_secs_f64() * 1000_f64; // time in ms
160
161    let now = Instant::now();
162    tree.apply_mutations(mutations)?;
163    let apply_elapsed = now.elapsed().as_secs_f64() * 1000_f64; // time in ms
164
165    println!(
166        "An average update-batch computation time measured by a {NUM_UPDATES}-batch into an SMT with {size} leaves over {:.1} ms is {:.0} μs",
167        compute_elapsed,
168        compute_elapsed * 1000_f64 / NUM_UPDATES as f64, // time in μs
169    );
170
171    println!(
172        "An average update-batch application time measured by a {NUM_UPDATES}-batch into an SMT with {size} leaves over {:.1} ms is {:.0} μs",
173        apply_elapsed,
174        apply_elapsed * 1000_f64 / NUM_UPDATES as f64, // time in μs
175    );
176
177    println!(
178        "An average batch update time measured by a 1k-batch into an SMT with {size} leaves totals to {:.1} ms",
179        (compute_elapsed + apply_elapsed),
180    );
181
182    println!();
183
184    Ok(())
185}
186
187/// Runs the proof generation benchmark for the [`Smt`].
188pub fn proof_generation(tree: &mut Smt) -> Result<(), MerkleError> {
189    const NUM_PROOFS: usize = 100;
190
191    println!("Running a proof generation benchmark:");
192
193    let mut insertion_times = Vec::new();
194
195    let size = tree.num_leaves();
196
197    for i in 0..NUM_PROOFS {
198        let test_key = Rpo256::hash(&rand_value::<u64>().to_be_bytes());
199        let test_value = [ONE, ONE, ONE, Felt::new((size + i) as u64)];
200        tree.insert(test_key, test_value);
201
202        let now = Instant::now();
203        let _proof = tree.open(&test_key);
204        insertion_times.push(now.elapsed().as_micros());
205    }
206
207    println!(
208        "An average proving time measured by {NUM_PROOFS} value proofs in an SMT with {size} leaves in {:.0} μs",
209        // calculate the average
210        insertion_times.iter().sum::<u128>() as f64 / (NUM_PROOFS as f64),
211    );
212
213    Ok(())
214}