Skip to main content

rupg_types/
prng.rs

1//! [`Prng`], the generator of `pg_prng.c`: xoroshiro128** with a seed from splitmix64. `random`, `setseed` and the random functions of the other types use it, so a session after `setseed` gives the same numbers as PostgreSQL.
2
3/// The state of the generator, as `pg_prng_state` holds it.
4#[derive(Debug, Clone, Copy, PartialEq, Eq)]
5pub struct Prng {
6    s0: u64,
7    s1: u64,
8}
9
10/// `splitmix64`: the next value of the generator of the seed.
11fn splitmix64(state: &mut u64) -> u64 {
12    *state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
13    let mut val = *state;
14    val = (val ^ (val >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
15    val = (val ^ (val >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
16    val ^ (val >> 31)
17}
18
19impl Prng {
20    /// `pg_prng_seed`: the state of a seed of 64 bits.
21    pub fn seed(seed: u64) -> Prng {
22        let mut seed = seed;
23        let s0 = splitmix64(&mut seed);
24        let s1 = splitmix64(&mut seed);
25        Prng::checked(s0, s1)
26    }
27
28    /// `pg_prng_strong_seed`: the state of 16 random bytes.
29    pub fn from_bytes(bytes: [u8; 16]) -> Prng {
30        let (s0, s1) = bytes.split_at(8);
31        let word = |b: &[u8]| u64::from_ne_bytes(b.try_into().unwrap_or_default());
32        Prng::checked(word(s0), word(s1))
33    }
34
35    /// `pg_prng_fseed`: the state of a seed in the range -1 to 1, as `setseed` gives it.
36    #[allow(clippy::cast_possible_truncation, clippy::cast_precision_loss, clippy::cast_sign_loss)]
37    pub fn fseed(fseed: f64) -> Prng {
38        let seed = (((1u64 << 52) - 1) as f64 * fseed) as i64;
39        Prng::seed(seed as u64)
40    }
41
42    /// `pg_prng_seed_check`: a state of all zeros becomes the parameters of the LCG of Knuth.
43    fn checked(s0: u64, s1: u64) -> Prng {
44        if s0 == 0 && s1 == 0 {
45            return Prng { s0: 0x5851_F42D_4C95_7F2D, s1: 0x1405_7B7E_F767_814F };
46        }
47        Prng { s0, s1 }
48    }
49
50    /// `xoroshiro128ss`: the next value, `pg_prng_uint64`.
51    pub fn next_u64(&mut self) -> u64 {
52        let s0 = self.s0;
53        let sx = self.s1 ^ s0;
54        let val = s0.wrapping_mul(5).rotate_left(7).wrapping_mul(9);
55        self.s0 = s0.rotate_left(24) ^ sx ^ (sx << 16);
56        self.s1 = sx.rotate_left(37);
57        val
58    }
59
60    /// `pg_prng_uint64_range`: a value from `rmin` to `rmax`, with the bitmask rejection method.
61    pub fn uint64_range(&mut self, rmin: u64, rmax: u64) -> u64 {
62        if rmax <= rmin {
63            return rmin;
64        }
65        let range = rmax - rmin;
66        let shift = range.leading_zeros();
67        loop {
68            let val = self.next_u64() >> shift;
69            if val <= range {
70                return rmin + val;
71            }
72        }
73    }
74
75    /// `pg_prng_int64_range`: a value from `rmin` to `rmax`.
76    #[allow(clippy::cast_sign_loss, clippy::cast_possible_wrap)]
77    pub fn int64_range(&mut self, rmin: i64, rmax: i64) -> i64 {
78        if rmax <= rmin {
79            return rmin;
80        }
81        let offset = self.uint64_range(0, (rmax as u64).wrapping_sub(rmin as u64));
82        (rmin as u64).wrapping_add(offset) as i64
83    }
84
85    /// `pg_prng_double`: a value from 0 up to but not including 1, with 52 random bits.
86    #[allow(clippy::cast_precision_loss)]
87    pub fn double(&mut self) -> f64 {
88        let v = self.next_u64() >> (64 - 52);
89        v as f64 / (1u64 << 52) as f64
90    }
91
92    /// `pg_prng_double_normal`: a value of the standard normal distribution, with the basic Box-Muller transform.
93    pub fn double_normal(&mut self) -> f64 {
94        let u1 = 1.0 - self.double();
95        let u2 = 1.0 - self.double();
96        (-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).sin()
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103
104    #[test]
105    fn range_is_inclusive() {
106        let mut prng = Prng::seed(1);
107        for _ in 0..1000 {
108            let v = prng.int64_range(-3, 3);
109            assert!((-3..=3).contains(&v));
110        }
111        assert_eq!(prng.int64_range(5, 5), 5);
112        assert_eq!(prng.uint64_range(0, 0), 0);
113    }
114
115    #[test]
116    fn double_is_below_one() {
117        let mut prng = Prng::fseed(0.5);
118        for _ in 0..1000 {
119            let v = prng.double();
120            assert!((0.0..1.0).contains(&v));
121        }
122    }
123
124    #[test]
125    fn zero_state_is_replaced() {
126        assert_eq!(Prng::from_bytes([0; 16]), Prng::checked(0, 0));
127        assert_ne!(Prng::checked(0, 0), Prng { s0: 0, s1: 0 });
128    }
129}