state.rs

2.4 kB · rust · 84 lines

1use crate::chacha::ChaCha8;2use std::collections::hash_map::RandomState;3use std::hash::{BuildHasher, Hasher};4use std::sync::{Mutex, OnceLock};56fn entropy() -> u64 {7    RandomState::new().build_hasher().finish()8}910fn rng() -> &'static Mutex<ChaCha8> {11    static RNG: OnceLock<Mutex<ChaCha8>> = OnceLock::new();12    RNG.get_or_init(|| Mutex::new(ChaCha8::from_u64(entropy())))13}1415/// Reseeds the global stream so every draw after it replays.16pub fn seed(s: u64) {17    *rng().lock().unwrap() = ChaCha8::from_u64(s);18}1920/// Draws a float at or above zero and below one from the global stream.21pub fn random() -> f64 {22    rng().lock().unwrap().unit()23}2425/// Draws a fair coin flip from the global stream.26pub fn boolean() -> bool {27    rng().lock().unwrap().boolean()28}2930/// Draws an integer between a and b inclusive from the global stream.31pub fn randint(a: i64, b: i64) -> i64 {32    rng().lock().unwrap().range_i64(a, b)33}3435/// Draws one element of the slice from the global stream.36pub fn choice<T: Clone>(seq: &[T]) -> T {37    let i = rng().lock().unwrap().below_u64(seq.len() as u64) as usize;38    seq[i].clone()39}4041/// Shuffles the slice in place with the global stream.42pub fn shuffle<T>(seq: &mut [T]) {43    rng().lock().unwrap().shuffle(seq);44}4546/// Draws k distinct elements of the slice from the global stream.47pub fn sample<T: Clone>(seq: &[T], k: usize) -> Vec<T> {48    let k = k.min(seq.len());49    let indices = rng().lock().unwrap().sample_indices(seq.len(), k);50    indices.into_iter().map(|i| seq[i].clone()).collect()51}5253/// Locks the global stream so one seeded test runs at a time.54#[cfg(any(test, feature = "testkit"))]55pub fn guard() -> std::sync::MutexGuard<'static, ()> {56    use std::sync::{Mutex, OnceLock};57    static LOCK: OnceLock<Mutex<()>> = OnceLock::new();58    LOCK.get_or_init(|| Mutex::new(()))59        .lock()60        .unwrap_or_else(|e| e.into_inner())61}6263#[cfg(test)]64mod tests {65    use super::*;66    #[test]67    fn seeded_is_reproducible() {68        let _g = guard();69        seed(42);70        let a: Vec<i64> = (0..10).map(|_| randint(0, 100)).collect();71        seed(42);72        let b: Vec<i64> = (0..10).map(|_| randint(0, 100)).collect();73        assert_eq!(a, b);74    }75    #[test]76    fn randint_is_inclusive() {77        let _g = guard();78        seed(7);79        for _ in 0..100 {80            let v = randint(1, 3);81            assert!((1..=3).contains(&v));82        }83    }84}