transaction_fuzzer/account_universe/
helpers.rs

1// Copyright (c) Mysten Labs, Inc.
2// Copyright (c) The Diem Core Contributors
3// Modifications Copyright (c) 2024 IOTA Stiftung
4// SPDX-License-Identifier: Apache-2.0
5
6use std::{
7    collections::BTreeSet,
8    ops::{Deref, Index as OpsIndex},
9};
10
11use proptest::sample::Index as PropIndex;
12use proptest_derive::Arbitrary;
13
14/// Given a maximum value `max` and a list of [`Index`](proptest::sample::Index)
15/// instances, picks integers in the range `[0, max)` uniformly randomly and
16/// without duplication.
17///
18/// If `indexes_len` is greater than `max`, all indexes will be returned.
19///
20/// This function implements [Robert Floyd's F2
21/// algorithm](https://blog.acolyer.org/2018/01/30/a-sample-of-brilliance/) for sampling without
22/// replacement.
23pub fn pick_idxs<T, P>(max: usize, indexes: &T, indexes_len: usize) -> Vec<usize>
24where
25    T: OpsIndex<usize, Output = P> + ?Sized,
26    P: AsRef<PropIndex>,
27{
28    // See https://blog.acolyer.org/2018/01/30/a-sample-of-brilliance/ (the F2 algorithm)
29    // for a longer explanation. This is a variant that works with zero-indexing.
30    let mut selected = BTreeSet::new();
31    let to_select = indexes_len.min(max);
32    for (iter_idx, choice) in ((max - to_select)..max).enumerate() {
33        // "RandInt(1, J)" in the original algorithm means a number between 1
34        // and choice, both inclusive. `PropIndex::index` picks a number between 0 and
35        // whatever's passed in, with the latter exclusive. Pass in "+1" to ensure the
36        // same range of values is picked from. (This also ensures that if
37        // choice is 0 then `index` doesn't panic.
38        let idx = indexes[iter_idx].as_ref().index(choice + 1);
39        if !selected.insert(idx) {
40            selected.insert(choice);
41        }
42    }
43    selected.into_iter().collect()
44}
45
46/// Given a maximum value `max` and a slice of
47/// [`Index`](proptest::sample::Index) instances, picks integers in the range
48/// `[0, max)` uniformly randomly and without duplication.
49///
50/// If the number of `Index` instances is greater than `max`, all indexes will
51/// be returned.
52///
53/// This function implements [Robert Floyd's F2
54/// algorithm](https://blog.acolyer.org/2018/01/30/a-sample-of-brilliance/) for sampling without
55/// replacement.
56#[inline]
57pub fn pick_slice_idxs(max: usize, indexes: &[impl AsRef<PropIndex>]) -> Vec<usize> {
58    pick_idxs(max, indexes, indexes.len())
59}
60/// Wrapper for `proptest`'s [`Index`][proptest::sample::Index] that allows
61/// `AsRef` to work.
62///
63/// There is no blanket `impl<T> AsRef<T> for T`, so `&[PropIndex]` doesn't work
64/// with `&[impl AsRef<PropIndex>]` (unless an impl gets added upstream).
65/// `Index` does.
66#[derive(Arbitrary, Clone, Copy, Debug)]
67pub struct Index(PropIndex);
68
69impl AsRef<PropIndex> for Index {
70    fn as_ref(&self) -> &PropIndex {
71        &self.0
72    }
73}
74
75impl Deref for Index {
76    type Target = PropIndex;
77
78    fn deref(&self) -> &PropIndex {
79        &self.0
80    }
81}