work_steal_queue/
rand.rs

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

work_steal_queue/
rand.rs
