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