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) -> ([u8; 32], Self) {
18 let mut seed = [0; 32];
19 self.source.fill_bytes(&mut seed);
20 (seed, Source::new(seed))
21 }
22
23 #[inline(always)]
24 pub fn next_u64n(&mut self, max: u64, mask: u64) -> u64 {
25 let mut x: u64 = self.next_u64() & mask;
26 while x >= max {
27 x = self.next_u64() & mask;
28 }
29 x
30 }
31
32 #[inline(always)]
33 pub fn next_f64(&mut self, min: f64, max: f64) -> f64 {
34 min + ((self.next_u64() << 11 >> 11) as f64) / MAXF64 * (max - min)
35 }
36
37 pub fn next_i64(&mut self) -> i64 {
38 self.next_u64() as i64
39 }
40}
41
42impl RngCore for Source {
43 #[inline(always)]
44 fn next_u32(&mut self) -> u32 {
45 self.source.next_u32()
46 }
47
48 #[inline(always)]
49 fn next_u64(&mut self) -> u64 {
50 self.source.next_u64()
51 }
52
53 #[inline(always)]
54 fn fill_bytes(&mut self, bytes: &mut [u8]) {
55 self.source.fill_bytes(bytes)
56 }
57}