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}