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