uno_core/quant/directms1/
rng.rs

1//! CPython integer-seeded MT19937 and _randbelow used by random.choice.
2//! This fixes resampling streams across Python and Rust, rather than substituting a different RNG.
3pub struct PythonRandom {
4    state: [u32; 624],
5    index: usize,
6}
7impl PythonRandom {
8    pub fn new(seed: u32) -> Self {
9        let mut r = Self {
10            state: [0; 624],
11            index: 624,
12        };
13        r.state[0] = 19650218;
14        for i in 1..624 {
15            r.state[i] = 1812433253u32
16                .wrapping_mul(r.state[i - 1] ^ (r.state[i - 1] >> 30))
17                .wrapping_add(i as u32);
18        }
19        let mut i = 1;
20        for _ in 0..624 {
21            r.state[i] = (r.state[i]
22                ^ (r.state[i - 1] ^ (r.state[i - 1] >> 30)).wrapping_mul(1664525))
23            .wrapping_add(seed);
24            i += 1;
25            if i >= 624 {
26                r.state[0] = r.state[623];
27                i = 1;
28            }
29        }
30        for _ in 0..623 {
31            r.state[i] = (r.state[i]
32                ^ (r.state[i - 1] ^ (r.state[i - 1] >> 30)).wrapping_mul(1566083941))
33            .wrapping_sub(i as u32);
34            i += 1;
35            if i >= 624 {
36                r.state[0] = r.state[623];
37                i = 1;
38            }
39        }
40        r.state[0] = 0x80000000;
41        r
42    }
43    fn next(&mut self) -> u32 {
44        if self.index == 624 {
45            for i in 0..624 {
46                let y = (self.state[i] & 0x80000000) | (self.state[(i + 1) % 624] & 0x7fffffff);
47                self.state[i] = self.state[(i + 397) % 624]
48                    ^ (y >> 1)
49                    ^ if y & 1 != 0 { 0x9908b0df } else { 0 };
50            }
51            self.index = 0;
52        }
53        let mut y = self.state[self.index];
54        self.index += 1;
55        y ^= y >> 11;
56        y ^= (y << 7) & 0x9d2c5680;
57        y ^= (y << 15) & 0xefc60000;
58        y ^= y >> 18;
59        y
60    }
61    pub fn choice(&mut self, n: usize) -> usize {
62        assert!(n > 0 && n <= u32::MAX as usize);
63        let bits = usize::BITS - n.leading_zeros();
64        loop {
65            let value = (self.next() >> (32 - bits)) as usize;
66            if value < n {
67                return value;
68            }
69        }
70    }
71}
72#[cfg(test)]
73mod tests {
74    use super::*;
75    #[test]
76    fn python_stream() {
77        let mut r = PythonRandom::new(42);
78        assert_eq!(
79            (0..10).map(|_| r.choice(1000)).collect::<Vec<_>>(),
80            vec![654, 114, 25, 759, 281, 250, 228, 142, 754, 104]
81        );
82    }
83}
uno_core/quant/directms1/rng.rs

uno_core/quant/directms1/
rng.rs
