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