rand_python/
lib.rs

1mod mersenne_twister;
2
3use std::slice;
4
5pub use mersenne_twister::MersenneTwister;
6
7type BigInt = u64; // could use `num` module to make this generic so you can use bigints etc
8
9
10pub struct PythonRandom {
11    mersenne_twister: MersenneTwister,
12}
13
14
15impl PythonRandom {
16    pub fn new(mt: MersenneTwister) -> PythonRandom {
17        PythonRandom {
18            mersenne_twister: mt,
19        }
20    }
21
22    pub fn random(&mut self) -> f64 {
23        self.mersenne_twister.genrand_res53()
24    }
25
26    pub fn getrandbits(&mut self, k: u32) -> BigInt {
27        assert!(0 < k && k <= BigInt::from(0u8).leading_zeros());
28
29        if k <= 32 {
30            return BigInt::from(self.mersenne_twister.genrand_int32() >> (32 - k));
31        }
32
33        let mut tmp = BigInt::from(0u8);
34        let mut k = k;
35        let mut shift = 0;
36        while k > 0 {
37            tmp |= BigInt::from(self.mersenne_twister.genrand_int32() >> 32u32.saturating_sub(k)) << shift;
38            k = k.saturating_sub(32);
39            shift += 32;
40        }
41
42        tmp
43    }
44
45    pub fn randbelow(&mut self, n: BigInt) -> BigInt {
46        let n_bits = BigInt::from(0u8).leading_zeros() - n.leading_zeros();
47        let mut r = self.getrandbits(n_bits);
48        while r >= n {
49            r = self.getrandbits(n_bits);
50        }
51        r
52    }
53
54    pub fn seed_u32(&mut self, s: u32) {
55        self.mersenne_twister.init_by_array(slice::from_ref(&s));
56    }
57
58    pub fn expovariate(&mut self, lambda: f64) -> f64 {
59        -(1.0 - self.mersenne_twister.genrand_res53()).ln() / lambda
60    }
61
62    pub fn shuffle<T>(&mut self, x: &mut [T]) {
63        for i in (1..x.len()).rev() {
64            let j = self.randbelow(BigInt::from(i as u64) + 1) as usize;
65            x.swap(i, j);
66        }
67    }
68
69    pub fn randint(&mut self, start: BigInt, stop: BigInt) -> BigInt {
70        self.randrange(start, stop + 1)
71    }
72
73    pub fn randrange(&mut self, start: BigInt, stop: BigInt) -> BigInt {
74        start + self.randbelow(stop - start)
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81
82    #[test]
83    fn sanity_checks() {
84        // Known-good values generated using the following python program:
85        //
86        // import random
87        // random.seed(63245986)
88        // print(random.getstate())
89        // print(random.random())
90        // print(random.randrange(0, 100000))
91        // print(random.expovariate(1.0 / 15000.0))
92        //
93        // tmp = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
94        // random.shuffle(tmp)
95        //
96        // print(tmp)
97
98        let mt = MersenneTwister::new();
99
100        let mut rand = PythonRandom::new(mt);
101
102        rand.seed_u32(63245986);
103
104        // println!("{:?}", &rand);
105
106        assert_eq!(rand.random(), 0.5213761361171212);
107
108        assert_eq!(rand.randbelow(100000u64), 58671);
109
110        assert_eq!(rand.expovariate(1.0 / 15000.0), 13775.46713470634);
111
112        let mut list: [u64; 10] = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
113
114        rand.shuffle(&mut list[..]);
115
116        assert_eq!(&list, &[10, 3, 6, 1, 8, 5, 7, 4, 2, 9]);
117    }
118}
rand_python/lib.rs

rand_python/
lib.rs
