arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/sampling/src/source.rs
sampling/src/source.rsRaw
 1use rand_chacha::{ChaCha8Rng, rand_core::SeedableRng};
 2use rand_core::RngCore;
 3
 4const MAXF64: f64 = 9007199254740992.0;
 5
 6pub struct Source {
 7    source: ChaCha8Rng,
 8}
 9
10impl Source {
11    pub fn new(seed: [u8; 32]) -> Source {
12        Source {
13            source: ChaCha8Rng::from_seed(seed),
14        }
15    }
16
17    pub fn branch(&mut self) -> Self {
18        let mut seed = [0; 32];
19        self.source.fill_bytes(&mut seed);
20
21        Source::new(seed)
22    }
23
24    #[inline(always)]
25    pub fn next_u64n(&mut self, max: u64, mask: u64) -> u64 {
26        let mut x: u64 = self.next_u64() & mask;
27        while x >= max {
28            x = self.next_u64() & mask;
29        }
30        x
31    }
32
33    #[inline(always)]
34    pub fn next_f64(&mut self, min: f64, max: f64) -> f64 {
35        min + ((self.next_u64() << 11 >> 11) as f64) / MAXF64 * (max - min)
36    }
37
38    pub fn next_i64(&mut self) -> i64 {
39        self.next_u64() as i64
40    }
41}
42
43impl RngCore for Source {
44    #[inline(always)]
45    fn next_u32(&mut self) -> u32 {
46        self.source.next_u32()
47    }
48
49    #[inline(always)]
50    fn next_u64(&mut self) -> u64 {
51        self.source.next_u64()
52    }
53
54    #[inline(always)]
55    fn fill_bytes(&mut self, bytes: &mut [u8]) {
56        self.source.fill_bytes(bytes)
57    }
58}