1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

//! Random module.

use std::cell::RefCell;

use num_bigint::{BigInt, Sign};
use num_traits::Signed;
use rand::RngCore;

use crate::function::OptionalOption;
use crate::obj::objint::PyIntRef;
use crate::obj::objtype::PyClassRef;
use crate::pyobject::{PyClassImpl, PyObjectRef, PyRef, PyResult, PyValue};
use crate::VirtualMachine;

mod mersenne;

#[derive(Debug)]
enum PyRng {
    Std(rand::rngs::ThreadRng),
    MT(Box<mersenne::MT19937>),
}

impl Default for PyRng {
    fn default() -> Self {
        PyRng::Std(rand::thread_rng())
    }
}

impl RngCore for PyRng {
    fn next_u32(&mut self) -> u32 {
        match self {
            Self::Std(s) => s.next_u32(),
            Self::MT(m) => m.next_u32(),
        }
    }
    fn next_u64(&mut self) -> u64 {
        match self {
            Self::Std(s) => s.next_u64(),
            Self::MT(m) => m.next_u64(),
        }
    }
    fn fill_bytes(&mut self, dest: &mut [u8]) {
        match self {
            Self::Std(s) => s.fill_bytes(dest),
            Self::MT(m) => m.fill_bytes(dest),
        }
    }
    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand::Error> {
        match self {
            Self::Std(s) => s.try_fill_bytes(dest),
            Self::MT(m) => m.try_fill_bytes(dest),
        }
    }
}

#[pyclass(name = "Random")]
#[derive(Debug)]
struct PyRandom {
    rng: RefCell<PyRng>,
}

impl PyValue for PyRandom {
    fn class(vm: &VirtualMachine) -> PyClassRef {
        vm.class("_random", "Random")
    }
}

#[pyimpl(flags(BASETYPE))]
impl PyRandom {
    #[pyslot(new)]
    fn new(cls: PyClassRef, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
        PyRandom {
            rng: RefCell::new(PyRng::default()),
        }
        .into_ref_with_type(vm, cls)
    }

    #[pymethod]
    fn random(&self) -> f64 {
        mersenne::gen_res53(&mut *self.rng.borrow_mut())
    }

    #[pymethod]
    fn seed(&self, n: OptionalOption<PyIntRef>) {
        let new_rng = match n.flat_option() {
            None => PyRng::default(),
            Some(n) => {
                let (_, mut key) = n.as_bigint().abs().to_u32_digits();
                if cfg!(target_endian = "big") {
                    key.reverse();
                }
                PyRng::MT(Box::new(mersenne::MT19937::new_with_slice_seed(&key)))
            }
        };

        *self.rng.borrow_mut() = new_rng;
    }

    #[pymethod]
    fn getrandbits(&self, mut k: usize) -> BigInt {
        let mut rng = self.rng.borrow_mut();

        let mut gen_u32 = |k| rng.next_u32() >> (32 - k) as u32;

        if k <= 32 {
            return gen_u32(k).into();
        }

        let words = (k - 1) / 8 + 1;
        let mut wordarray = vec![0u32; words];

        let it = wordarray.iter_mut();
        #[cfg(target_endian = "big")]
        let it = it.rev();
        for word in it {
            *word = gen_u32(k);
            k -= 32;
        }

        BigInt::from_slice(Sign::NoSign, &wordarray)
    }
}

pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
    let ctx = &vm.ctx;
    py_module!(vm, "_random", {
        "Random" => PyRandom::make_class(ctx),
    })
}