token-sampling 0.1.0

/// Random bits used by sampling stages.
pub trait Rng {
    /// Returns the next 64 random bits.
    fn next_u64(&mut self) -> u64;
}

/// Converts one random word to a value in `[0, 1)` with 53 bits of precision.
#[inline]
pub fn unit(rng: &mut dyn Rng) -> f64 {
    ((rng.next_u64() >> 11) as f64) * (1.0 / 9_007_199_254_740_992.0)
}

/// A small seeded generator based on SplitMix64.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct SplitMix64 {
    state: u64,
}

impl SplitMix64 {
    /// Creates a generator from `seed`.
    pub const fn new(seed: u64) -> Self {
        Self { state: seed }
    }
}

impl Rng for SplitMix64 {
    #[inline]
    fn next_u64(&mut self) -> u64 {
        self.state = self.state.wrapping_add(0x9e37_79b9_7f4a_7c15);
        let mut value = self.state;
        value = (value ^ (value >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
        value = (value ^ (value >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
        value ^ (value >> 31)
    }
}

/// An RNG that replays borrowed random words.
///
/// Exhausted tapes return zero. This keeps tests bounded when a stage consumes
/// more draws than expected.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TapeRng<'a> {
    tape: &'a [u64],
    position: usize,
}

impl<'a> TapeRng<'a> {
    /// Creates a tape at its first word.
    pub const fn new(tape: &'a [u64]) -> Self {
        Self { tape, position: 0 }
    }

    /// Returns the number of unread words.
    pub fn remaining(&self) -> usize {
        self.tape.len().saturating_sub(self.position)
    }
}

impl Rng for TapeRng<'_> {
    fn next_u64(&mut self) -> u64 {
        let value = self.tape.get(self.position).copied().unwrap_or(0);
        self.position = self.position.saturating_add(1);
        value
    }
}

#[cfg(feature = "rand_core")]
impl<T: rand_core::RngCore + ?Sized> Rng for T {
    #[inline]
    fn next_u64(&mut self) -> u64 {
        rand_core::RngCore::next_u64(self)
    }
}