token_sampling/
rng.rs

1/// Random bits used by sampling stages.
2pub trait Rng {
3    /// Returns the next 64 random bits.
4    fn next_u64(&mut self) -> u64;
5}
6
7/// Converts one random word to a value in `[0, 1)` with 53 bits of precision.
8#[inline]
9pub fn unit(rng: &mut dyn Rng) -> f64 {
10    ((rng.next_u64() >> 11) as f64) * (1.0 / 9_007_199_254_740_992.0)
11}
12
13/// A small seeded generator based on SplitMix64.
14#[derive(Clone, Copy, Debug, PartialEq, Eq)]
15pub struct SplitMix64 {
16    state: u64,
17}
18
19impl SplitMix64 {
20    /// Creates a generator from `seed`.
21    pub const fn new(seed: u64) -> Self {
22        Self { state: seed }
23    }
24}
25
26impl Rng for SplitMix64 {
27    #[inline]
28    fn next_u64(&mut self) -> u64 {
29        self.state = self.state.wrapping_add(0x9e37_79b9_7f4a_7c15);
30        let mut value = self.state;
31        value = (value ^ (value >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
32        value = (value ^ (value >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
33        value ^ (value >> 31)
34    }
35}
36
37/// An RNG that replays borrowed random words.
38///
39/// Exhausted tapes return zero. This keeps tests bounded when a stage consumes
40/// more draws than expected.
41#[derive(Clone, Debug, PartialEq, Eq)]
42pub struct TapeRng<'a> {
43    tape: &'a [u64],
44    position: usize,
45}
46
47impl<'a> TapeRng<'a> {
48    /// Creates a tape at its first word.
49    pub const fn new(tape: &'a [u64]) -> Self {
50        Self { tape, position: 0 }
51    }
52
53    /// Returns the number of unread words.
54    pub fn remaining(&self) -> usize {
55        self.tape.len().saturating_sub(self.position)
56    }
57}
58
59impl Rng for TapeRng<'_> {
60    fn next_u64(&mut self) -> u64 {
61        let value = self.tape.get(self.position).copied().unwrap_or(0);
62        self.position = self.position.saturating_add(1);
63        value
64    }
65}
66
67#[cfg(feature = "rand_core")]
68impl<T: rand_core::RngCore + ?Sized> Rng for T {
69    #[inline]
70    fn next_u64(&mut self) -> u64 {
71        rand_core::RngCore::next_u64(self)
72    }
73}
token_sampling/rng.rs

token_sampling/
rng.rs
