Skip to main content

mecha_core/
sample.rs

1//! Reproducible random draws.
2//!
3//! Two places already needed one and each wrote its own: the mail corpus
4//! grader, and the graph's queue sampler one repository over. This is the
5//! third caller — prioritised replay's holdout — and a third copy of a
6//! shuffle is a third place for the same bias to hide.
7//!
8//! **The seed is the caller's and gets printed.** A sample nobody can redraw
9//! is a sample nobody can check, and these exist to produce numbers somebody
10//! will quote.
11//!
12//! **Sort before you shuffle, or the seed is a lie.** The mail grader learned
13//! this the expensive way: it shuffled a `HashMap`'s iteration order, which is
14//! randomised per process, so two runs with the same seed graded different
15//! samples while the flag documented itself as making a scorecard
16//! reproducible. A deterministic shuffle of a nondeterministic order is
17//! nondeterministic. Callers pass a slice whose order they control.
18//!
19//! The PRNG is a four-line LCG rather than a dependency: nothing here needs
20//! cryptographic randomness, and `rand` would be a new dependency in a crate
21//! that reads the owner's transcripts.
22
23/// Fisher–Yates, seeded.
24///
25/// Backward, and over the whole vector — [`take_uniform`] then takes a prefix,
26/// which is uniform because the shuffle was complete. Truncating an *unshuffled*
27/// list is the bias this exists to escape.
28pub fn shuffled<T>(mut v: Vec<T>, seed: u64) -> Vec<T> {
29    let mut state = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
30    let mut next = || {
31        state = state
32            .wrapping_mul(6364136223846793005)
33            .wrapping_add(1442695040888963407);
34        (state >> 33) as usize
35    };
36    for i in (1..v.len()).rev() {
37        v.swap(i, next() % (i + 1));
38    }
39    v
40}
41
42/// A uniform draw of at most `k`, reproducible from `seed`.
43pub fn take_uniform<T>(v: Vec<T>, seed: u64, k: usize) -> Vec<T> {
44    shuffled(v, seed).into_iter().take(k).collect()
45}
46
47#[cfg(test)]
48mod tests {
49    use super::*;
50
51    #[test]
52    fn the_same_seed_draws_the_same_sample_and_a_different_one_does_not() {
53        let items: Vec<u32> = (0..40).collect();
54        assert_eq!(
55            take_uniform(items.clone(), 42, 8),
56            take_uniform(items.clone(), 42, 8)
57        );
58        assert_ne!(
59            take_uniform(items.clone(), 42, 8),
60            take_uniform(items, 43, 8)
61        );
62    }
63
64    /// The property the draw exists for. Every item must appear in the drawn
65    /// prefix at close to the same rate — this is what fails when a caller
66    /// reaches for `truncate(k)` on an unshuffled list, which is the shape the
67    /// graph's sampler had to escape and the reason this is not open-coded a
68    /// third time.
69    #[test]
70    fn every_item_is_drawn_at_close_to_the_same_rate() {
71        const N: usize = 10;
72        const K: usize = 3;
73        const DRAWS: u64 = 4_000;
74        let mut seen = [0usize; N];
75        for seed in 0..DRAWS {
76            for i in take_uniform((0..N).collect::<Vec<_>>(), seed, K) {
77                seen[i] += 1;
78            }
79        }
80        let expected = (DRAWS as f64) * (K as f64) / (N as f64);
81        for (i, count) in seen.iter().enumerate() {
82            let drift = (*count as f64 - expected).abs() / expected;
83            assert!(
84                drift < 0.15,
85                "item {i} drawn {count} times against an expected {expected:.0} ({:.1}% off)",
86                drift * 100.0
87            );
88        }
89    }
90
91    #[test]
92    fn a_draw_larger_than_the_pool_returns_the_pool() {
93        assert_eq!(take_uniform(vec![1, 2, 3], 7, 99).len(), 3);
94        assert!(take_uniform(Vec::<u8>::new(), 7, 5).is_empty());
95    }
96}