rng.rs

1use crate::chacha::ChaCha8;23/// A seeded, seekable ChaCha8 random stream.4#[derive(Clone)]5pub struct Rng {6    inner: ChaCha8,7}89impl Rng {10    /// Builds the stream from a seed.11    pub fn new(seed: u64) -> Rng {12        Rng {13            inner: ChaCha8::from_u64(seed),14        }15    }16    /// Draws an integer between lo and hi inclusive, or lo when hi is not above lo.17    pub fn range(&mut self, lo: i64, hi: i64) -> i64 {18        if hi <= lo {19            lo20        } else {21            self.inner.range_i64(lo, hi)22        }23    }24    /// Draws an integer below n, or zero when n is zero.25    pub fn below(&mut self, n: usize) -> usize {26        if n == 0 {27            028        } else {29            self.inner.below_u64(n as u64) as usize30        }31    }32    /// Draws a float at or above zero and below one.33    pub fn unit(&mut self) -> f64 {34        self.inner.unit()35    }36    /// Draws a fair coin flip.37    pub fn boolean(&mut self) -> bool {38        self.inner.boolean()39    }40    /// Returns true with probability p.41    pub fn chance(&mut self, p: f64) -> bool {42        self.unit() < p43    }44    /// Draws one element of the slice.45    pub fn choice<'a, T>(&mut self, items: &'a [T]) -> &'a T {46        &items[self.below(items.len())]47    }48    /// Returns the stream's current word position.49    pub fn pos(&self) -> u128 {50        self.inner.word_pos()51    }52    /// Moves the stream to a word position so later draws replay from there.53    pub fn seek(&mut self, pos: u128) {54        self.inner.set_word_pos(pos);55    }56}5758#[cfg(test)]59mod tests {60    use super::*;61    #[test]62    fn seek_resumes_the_stream() {63        let mut a = Rng::new(7);64        a.below(100);65        a.unit();66        let pos = a.pos();67        let next: Vec<usize> = (0..5).map(|_| a.below(1000)).collect();68        let mut b = Rng::new(7);69        b.seek(pos);70        let again: Vec<usize> = (0..5).map(|_| b.below(1000)).collect();71        assert_eq!(next, again);72    }73}