monoio/utils/
rand.rs

1//! Fast random number generate
2//!
3//! Implement xorshift64+: 2 32-bit xorshift sequences added together.
4//! Shift triplet `[17,7,16]` was calculated as indicated in Marsaglia's
5//! Xorshift paper: <https://www.jstatsoft.org/article/view/v008i14/xorshift.pdf>
6//! This generator passes the SmallCrush suite, part of TestU01 framework:
7//! <http://simul.iro.umontreal.ca/testu01/tu01.html>
8// Heavily borrowed from tokio.
9// Copyright (c) 2021 Tokio Contributors, licensed under the MIT license.
10use std::cell::Cell;
11
12#[derive(Debug)]
13pub(crate) struct FastRand {
14    one: Cell<u32>,
15    two: Cell<u32>,
16}
17
18impl FastRand {
19    /// Initialize a new, thread-local, fast random number generator.
20    pub(crate) fn new(seed: u64) -> FastRand {
21        let one = (seed >> 32) as u32;
22        let mut two = seed as u32;
23
24        if two == 0 {
25            // This value cannot be zero
26            two = 1;
27        }
28
29        FastRand {
30            one: Cell::new(one),
31            two: Cell::new(two),
32        }
33    }
34
35    pub(crate) fn fastrand_n(&self, n: u32) -> u32 {
36        // This is similar to fastrand() % n, but faster.
37        // See https://lemire.me/blog/2016/06/27/a-fast-alternative-to-the-modulo-reduction/
38        let mul = (self.fastrand() as u64).wrapping_mul(n as u64);
39        (mul >> 32) as u32
40    }
41
42    fn fastrand(&self) -> u32 {
43        let mut s1 = self.one.get();
44        let s0 = self.two.get();
45
46        s1 ^= s1 << 17;
47        s1 = s1 ^ s0 ^ (s1 >> 7) ^ (s0 >> 16);
48
49        self.one.set(s0);
50        self.two.set(s1);
51
52        s0.wrapping_add(s1)
53    }
54}
55
56/// Used by the select macro and `StreamMap`
57pub fn thread_rng_n(n: u32) -> u32 {
58    thread_local! {
59        static THREAD_RNG: FastRand = FastRand::new(seed());
60    }
61
62    THREAD_RNG.with(|rng| rng.fastrand_n(n))
63}
64
65use std::{
66    collections::hash_map::RandomState,
67    hash::BuildHasher,
68    sync::atomic::{AtomicU32, Ordering::Relaxed},
69};
70
71static COUNTER: AtomicU32 = AtomicU32::new(1);
72
73fn seed() -> u64 {
74    let rand_state = RandomState::new();
75    rand_state.hash_one(COUNTER.fetch_add(1, Relaxed))
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81
82    #[test]
83    fn rand() {
84        for _ in 0..100 {
85            assert!(thread_rng_n(10) < 10);
86        }
87    }
88}
monoio/utils/rand.rs

monoio/utils/
rand.rs
