1use std::{fmt::Debug, hint, mem, time::Duration};
  2
  3use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
  4use miden_crypto::{
  5    hash::rpo::RpoDigest,
  6    merkle::{build_subtree_for_bench, NodeIndex, SmtLeaf, SubtreeLeaf, SMT_DEPTH},
  7    Felt, Word, ONE,
  8};
  9use rand_utils::prng_array;
 10use winter_utils::Randomizable;
 11
 12const PAIR_COUNTS: [u64; 5] = [1, 64, 128, 192, 256];
 13
 14fn smt_subtree_even(c: &mut Criterion) {
 15    let mut seed = [0u8; 32];
 16
 17    let mut group = c.benchmark_group("subtree8-even");
 18
 19    for pair_count in PAIR_COUNTS {
 20        let bench_id = BenchmarkId::from_parameter(pair_count);
 21        group.bench_with_input(bench_id, &pair_count, |b, &pair_count| {
 22            b.iter_batched(
 23                || {
 24                    // Setup.
 25                    let entries: Vec<(RpoDigest, Word)> = (0..pair_count)
 26                        .map(|n| {
 27                            // A single depth-8 subtree can have a maximum of 255 leaves.
 28                            let leaf_index = ((n as f64 / pair_count as f64) * 255.0) as u64;
 29                            let key = RpoDigest::new([
 30                                generate_value(&mut seed),
 31                                ONE,
 32                                Felt::new(n),
 33                                Felt::new(leaf_index),
 34                            ]);
 35                            let value = generate_word(&mut seed);
 36                            (key, value)
 37                        })
 38                        .collect();
 39
 40                    let mut leaves: Vec<_> = entries
 41                        .iter()
 42                        .map(|(key, value)| {
 43                            let leaf = SmtLeaf::new_single(*key, *value);
 44                            let col = NodeIndex::from(leaf.index()).value();
 45                            let hash = leaf.hash();
 46                            SubtreeLeaf { col, hash }
 47                        })
 48                        .collect();
 49                    leaves.sort();
 50                    leaves.dedup_by_key(|leaf| leaf.col);
 51                    leaves
 52                },
 53                |leaves| {
 54                    // Benchmarked function.
 55                    let (subtree, _) = build_subtree_for_bench(
 56                        hint::black_box(leaves),
 57                        hint::black_box(SMT_DEPTH),
 58                        hint::black_box(SMT_DEPTH),
 59                    );
 60                    assert!(!subtree.is_empty());
 61                },
 62                BatchSize::SmallInput,
 63            );
 64        });
 65    }
 66}
 67
 68fn smt_subtree_random(c: &mut Criterion) {
 69    let mut seed = [0u8; 32];
 70
 71    let mut group = c.benchmark_group("subtree8-rand");
 72
 73    for pair_count in PAIR_COUNTS {
 74        let bench_id = BenchmarkId::from_parameter(pair_count);
 75        group.bench_with_input(bench_id, &pair_count, |b, &pair_count| {
 76            b.iter_batched(
 77                || {
 78                    // Setup.
 79                    let entries: Vec<(RpoDigest, Word)> = (0..pair_count)
 80                        .map(|i| {
 81                            let leaf_index: u8 = generate_value(&mut seed);
 82                            let key = RpoDigest::new([
 83                                ONE,
 84                                ONE,
 85                                Felt::new(i),
 86                                Felt::new(leaf_index as u64),
 87                            ]);
 88                            let value = generate_word(&mut seed);
 89                            (key, value)
 90                        })
 91                        .collect();
 92
 93                    let mut leaves: Vec<_> = entries
 94                        .iter()
 95                        .map(|(key, value)| {
 96                            let leaf = SmtLeaf::new_single(*key, *value);
 97                            let col = NodeIndex::from(leaf.index()).value();
 98                            let hash = leaf.hash();
 99                            SubtreeLeaf { col, hash }
100                        })
101                        .collect();
102                    leaves.sort();
103                    leaves
104                },
105                |leaves| {
106                    let (subtree, _) = build_subtree_for_bench(
107                        hint::black_box(leaves),
108                        hint::black_box(SMT_DEPTH),
109                        hint::black_box(SMT_DEPTH),
110                    );
111                    assert!(!subtree.is_empty());
112                },
113                BatchSize::SmallInput,
114            );
115        });
116    }
117}
118
119criterion_group! {
120    name = smt_subtree_group;
121    config = Criterion::default()
122        .measurement_time(Duration::from_secs(40))
123        .sample_size(60)
124        .configure_from_args();
125    targets = smt_subtree_even, smt_subtree_random
126}
127criterion_main!(smt_subtree_group);
128
129// HELPER FUNCTIONS
130// --------------------------------------------------------------------------------------------
131
132fn generate_value<T: Copy + Debug + Randomizable>(seed: &mut [u8; 32]) -> T {
133    mem::swap(seed, &mut prng_array(*seed));
134    let value: [T; 1] = rand_utils::prng_array(*seed);
135    value[0]
136}
137
138fn generate_word(seed: &mut [u8; 32]) -> Word {
139    mem::swap(seed, &mut prng_array(*seed));
140    let nums: [u64; 4] = prng_array(*seed);
141    [Felt::new(nums[0]), Felt::new(nums[1]), Felt::new(nums[2]), Felt::new(nums[3])]
142}