open_coroutine_queue/
rand.rs

1use std::cell::Cell;
2
3use parking_lot::Mutex;
4use std::collections::hash_map::RandomState;
5use std::hash::BuildHasher;
6use std::sync::atomic::AtomicU32;
7use std::sync::atomic::Ordering::Relaxed;
8
9static COUNTER: AtomicU32 = AtomicU32::new(1);
10
11pub(crate) fn seed() -> u64 {
12    let rand_state = RandomState::new();
13    // Hash some unique-ish data to generate some new state
14    // Get the seed
15    rand_state.hash_one(COUNTER.fetch_add(1, Relaxed))
16}
17
18/// A deterministic generator for seeds (and other generators).
19///
20/// Given the same initial seed, the generator will output the same sequence of seeds.
21///
22/// Since the seed generator will be kept in a runtime handle, we need to wrap `FastRand`
23/// in a Mutex to make it thread safe. Different to the `FastRand` that we keep in a
24/// thread local store, the expectation is that seed generation will not need to happen
25/// very frequently, so the cost of the mutex should be minimal.
26#[repr(C)]
27#[derive(Debug)]
28pub struct RngSeedGenerator {
29    /// Internal state for the seed generator. We keep it in a Mutex so that we can safely
30    /// use it across multiple threads.
31    state: Mutex<FastRand>,
32}
33
34impl RngSeedGenerator {
35    /// Returns a new generator from the provided seed.
36    #[must_use]
37    pub fn new(seed: RngSeed) -> Self {
38        Self {
39            state: Mutex::new(FastRand::new(seed)),
40        }
41    }
42
43    /// Returns the next seed in the sequence.
44    pub fn next_seed(&self) -> RngSeed {
45        let rng = self.state.lock();
46
47        let s = rng.fastrand();
48        let r = rng.fastrand();
49
50        RngSeed::from_pair(s, r)
51    }
52
53    /// Directly creates a generator using the next seed.
54    #[must_use]
55    pub fn next_generator(&self) -> Self {
56        RngSeedGenerator::new(self.next_seed())
57    }
58}
59
60impl Default for RngSeedGenerator {
61    fn default() -> Self {
62        Self::new(RngSeed::new())
63    }
64}
65
66/// A seed for random number generation.
67///
68/// In order to make certain functions within a runtime deterministic, a seed
69/// can be specified at the time of creation.
70#[allow(unreachable_pub)]
71#[derive(Debug, Copy, Clone)]
72pub struct RngSeed {
73    s: u32,
74    r: u32,
75}
76
77impl RngSeed {
78    /// Creates a random seed using loom internally.
79    #[must_use]
80    pub fn new() -> Self {
81        Self::from_u64(seed())
82    }
83
84    #[allow(clippy::cast_possible_truncation)]
85    fn from_u64(seed: u64) -> Self {
86        let one = (seed >> 32) as u32;
87        let mut two = seed as u32;
88
89        if two == 0 {
90            // This value cannot be zero
91            two = 1;
92        }
93
94        Self::from_pair(one, two)
95    }
96
97    fn from_pair(s: u32, r: u32) -> Self {
98        Self { s, r }
99    }
100}
101
102impl Default for RngSeed {
103    fn default() -> Self {
104        Self::new()
105    }
106}
107
108/// Fast random number generate.
109///
110/// Implement xorshift64+: 2 32-bit xorshift sequences added together.
111/// Shift triplet `[17,7,16]` was calculated as indicated in Marsaglia's
112/// Xorshift paper: <https://www.jstatsoft.org/article/view/v008i14/xorshift.pdf>
113/// This generator passes the `SmallCrush` suite, part of `TestU01` framework:
114/// <http://simul.iro.umontreal.ca/testu01/tu01.html>
115#[repr(C)]
116#[derive(Debug)]
117pub struct FastRand {
118    one: Cell<u32>,
119    two: Cell<u32>,
120}
121
122impl FastRand {
123    /// Initializes a new, thread-local, fast random number generator.
124    #[must_use]
125    pub fn new(seed: RngSeed) -> FastRand {
126        FastRand {
127            one: Cell::new(seed.s),
128            two: Cell::new(seed.r),
129        }
130    }
131
132    /// Replaces the state of the random number generator with the provided seed, returning
133    /// the seed that represents the previous state of the random number generator.
134    ///
135    /// The random number generator will become equivalent to one created with
136    /// the same seed.
137    pub fn replace_seed(&self, seed: RngSeed) -> RngSeed {
138        let old_seed = RngSeed::from_pair(self.one.get(), self.two.get());
139
140        _ = self.one.replace(seed.s);
141        _ = self.two.replace(seed.r);
142
143        old_seed
144    }
145
146    pub fn fastrand_n(&self, n: u32) -> u32 {
147        // This is similar to fastrand() % n, but faster.
148        // See https://lemire.me/blog/2016/06/27/a-fast-alternative-to-the-modulo-reduction/
149        let mul = (u64::from(self.fastrand())).wrapping_mul(u64::from(n));
150        (mul >> 32) as u32
151    }
152
153    fn fastrand(&self) -> u32 {
154        let mut s1 = self.one.get();
155        let s0 = self.two.get();
156
157        s1 ^= s1 << 17;
158        s1 = s1 ^ s0 ^ s1 >> 7 ^ s0 >> 16;
159
160        self.one.set(s0);
161        self.two.set(s1);
162
163        s0.wrapping_add(s1)
164    }
165}