1use core::mem::swap;
 2
 3use criterion::{black_box, criterion_group, criterion_main, Criterion};
 4use miden_crypto::{
 5    merkle::{LeafIndex, SimpleSmt},
 6    Felt, Word,
 7};
 8use rand_utils::prng_array;
 9use seq_macro::seq;
10
11fn smt_rpo(c: &mut Criterion) {
12    // setup trees
13
14    let mut seed = [0u8; 32];
15    let leaf = generate_word(&mut seed);
16
17    seq!(DEPTH in 14..=20 {
18        let leaves = ((1 << DEPTH) - 1) as u64;
19        for count in [1, leaves / 2, leaves] {
20            let entries: Vec<_> = (0..count)
21                .map(|i| {
22                    let word = generate_word(&mut seed);
23                    (i, word)
24                })
25                .collect();
26            let mut tree = SimpleSmt::<DEPTH>::with_leaves(entries).unwrap();
27
28            // benchmark 1
29            let mut insert = c.benchmark_group("smt update_leaf".to_string());
30            {
31                let depth = DEPTH;
32                let key = count >> 2;
33                insert.bench_with_input(
34                    format!("simple smt(depth:{depth},count:{count})"),
35                    &(key, leaf),
36                    |b, (key, leaf)| {
37                        b.iter(|| {
38                            tree.insert(black_box(LeafIndex::<DEPTH>::new(*key).unwrap()), black_box(*leaf));
39                        });
40                    },
41                );
42
43            }
44            insert.finish();
45
46            // benchmark 2
47            let mut path = c.benchmark_group("smt get_leaf_path".to_string());
48            {
49                let depth = DEPTH;
50                let key = count >> 2;
51                path.bench_with_input(
52                    format!("simple smt(depth:{depth},count:{count})"),
53                    &key,
54                    |b, key| {
55                        b.iter(|| {
56                            tree.open(black_box(&LeafIndex::<DEPTH>::new(*key).unwrap()));
57                        });
58                    },
59                );
60
61            }
62            path.finish();
63        }
64    });
65}
66
67criterion_group!(smt_group, smt_rpo);
68criterion_main!(smt_group);
69
70// HELPER FUNCTIONS
71// --------------------------------------------------------------------------------------------
72
73fn generate_word(seed: &mut [u8; 32]) -> Word {
74    swap(seed, &mut prng_array(*seed));
75    let nums: [u64; 4] = prng_array(*seed);
76    [Felt::new(nums[0]), Felt::new(nums[1]), Felt::new(nums[2]), Felt::new(nums[3])]
77}