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::{OsRng, RngCore, TryRngCore};
 4
 5const MAXF64: f64 = 9007199254740992.0;
 6
 7pub struct Source {
 8    source: ChaCha8Rng,
 9}
10
11pub fn new_seed() -> [u8; 32] {
12    let mut seed = [0u8; 32];
13    OsRng.try_fill_bytes(&mut seed).unwrap();
14    seed
15}
16
17impl Source {
18    pub fn new(seed: [u8; 32]) -> Source {
19        Source {
20            source: ChaCha8Rng::from_seed(seed),
21        }
22    }
23
24    pub fn new_seed(&mut self) -> [u8; 32] {
25        let mut seed: [u8; 32] = [0u8; 32];
26        self.source.fill_bytes(&mut seed);
27        seed
28    }
29
30    pub fn branch(&mut self) -> Self {
31        Source::new(self.new_seed())
32    }
33
34    #[inline(always)]
35    pub fn next_u64n(&mut self, max: u64, mask: u64) -> u64 {
36        let mut x: u64 = self.next_u64() & mask;
37        while x >= max {
38            x = self.next_u64() & mask;
39        }
40        x
41    }
42
43    #[inline(always)]
44    pub fn next_f64(&mut self, min: f64, max: f64) -> f64 {
45        min + ((self.next_u64() << 11 >> 11) as f64) / MAXF64 * (max - min)
46    }
47
48    pub fn next_i64(&mut self) -> i64{
49        self.next_u64() as i64
50    }
51}
52
53impl RngCore for Source {
54    #[inline(always)]
55    fn next_u32(&mut self) -> u32 {
56        self.source.next_u32()
57    }
58
59    #[inline(always)]
60    fn next_u64(&mut self) -> u64 {
61        self.source.next_u64()
62    }
63
64    #[inline(always)]
65    fn fill_bytes(&mut self, bytes: &mut [u8]) {
66        self.source.fill_bytes(bytes)
67    }
68}