1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

// Copyright 2021 Conflux Foundation. All rights reserved.
// Conflux is free software and distributed under GNU General Public License.
// See http://www.gnu.org/licenses/

#![forbid(unsafe_code)]

#[cfg(test)]
mod unit_tests;

mod growing_subset;
mod repeat_vec;
mod value_generator;

pub use crate::{
    growing_subset::GrowingSubset, repeat_vec::RepeatVec,
    value_generator::ValueGenerator,
};

use crossbeam::thread;
use proptest::sample::Index as PropIndex;
use proptest_derive::Arbitrary;
use std::{
    any::Any,
    collections::BTreeSet,
    ops::{Deref, Index as OpsIndex},
};

/// Creates a new thread with a larger stack size.
///
/// Generating some proptest values can overflow the stack. This allows test
/// authors to work around this limitation.
///
/// This is expected to be used with closure-style proptest invocations:
///
/// ```
/// use proptest::prelude::*;
/// use diem_proptest_helpers::with_stack_size;
///
/// with_stack_size(4 * 1024 * 1024, || proptest!(|(x in 0usize..128)| {
///     // assertions go here
///     prop_assert!(x >= 0 && x < 128);
/// }));
/// ```
pub fn with_stack_size<'a, F, T>(
    size: usize, f: F,
) -> Result<T, Box<dyn Any + 'static + Send>>
where
    F: FnOnce() -> T + Send + 'a,
    T: Send + 'a,
{
    thread::scope(|s| {
        let handle =
            s.builder().stack_size(size).spawn(|_| f()).map_err(|err| {
                let any: Box<dyn Any + 'static + Send> = Box::new(err);
                any
            })?;
        handle.join()
    })?
}

/// Given a maximum value `max` and a list of [`Index`](proptest::sample::Index)
/// instances, picks integers in the range `[0, max)` uniformly randomly and
/// without duplication.
///
/// If `indexes_len` is greater than `max`, all indexes will be returned.
///
/// This function implements [Robert Floyd's F2
/// algorithm](https://blog.acolyer.org/2018/01/30/a-sample-of-brilliance/) for sampling without
/// replacement.
pub fn pick_idxs<T, P>(
    max: usize, indexes: &T, indexes_len: usize,
) -> Vec<usize>
where
    T: OpsIndex<usize, Output = P> + ?Sized,
    P: AsRef<PropIndex>,
{
    // See https://blog.acolyer.org/2018/01/30/a-sample-of-brilliance/ (the F2 algorithm)
    // for a longer explanation. This is a variant that works with
    // zero-indexing.
    let mut selected = BTreeSet::new();
    let to_select = indexes_len.min(max);
    for (iter_idx, choice) in ((max - to_select)..max).enumerate() {
        // "RandInt(1, J)" in the original algorithm means a number between 1
        // and choice, both inclusive. `PropIndex::index` picks a number between
        // 0 and whatever's passed in, with the latter exclusive. Pass
        // in "+1" to ensure the same range of values is picked from.
        // (This also ensures that if choice is 0 then `index`
        // doesn't panic.
        let idx = indexes[iter_idx].as_ref().index(choice + 1);
        if !selected.insert(idx) {
            selected.insert(choice);
        }
    }
    selected.into_iter().collect()
}

/// Given a maximum value `max` and a slice of
/// [`Index`](proptest::sample::Index) instances, picks integers in the range
/// `[0, max)` uniformly randomly and without duplication.
///
/// If the number of `Index` instances is greater than `max`, all indexes will
/// be returned.
///
/// This function implements [Robert Floyd's F2
/// algorithm](https://blog.acolyer.org/2018/01/30/a-sample-of-brilliance/) for sampling without
/// replacement.
#[inline]
pub fn pick_slice_idxs(
    max: usize, indexes: &[impl AsRef<PropIndex>],
) -> Vec<usize> {
    pick_idxs(max, indexes, indexes.len())
}

/// Wrapper for `proptest`'s [`Index`][proptest::sample::Index] that allows
/// `AsRef` to work.
///
/// There is no blanket `impl<T> AsRef<T> for T`, so `&[PropIndex]` doesn't work
/// with `&[impl AsRef<PropIndex>]` (unless an impl gets added upstream).
/// `Index` does.
#[derive(Arbitrary, Clone, Copy, Debug)]
pub struct Index(PropIndex);

impl AsRef<PropIndex> for Index {
    fn as_ref(&self) -> &PropIndex { &self.0 }
}

impl Deref for Index {
    type Target = PropIndex;

    fn deref(&self) -> &PropIndex { &self.0 }
}