diff --git a/Cargo.lock b/Cargo.lock index 023c6240c..66c957bc4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1575,6 +1575,7 @@ dependencies = [ "flamer", "flate2", "foreign-types-shared", + "generational-arena", "gethostname", "getrandom", "hex", diff --git a/vm/Cargo.toml b/vm/Cargo.toml index 71c179e3e..3c5dada10 100644 --- a/vm/Cargo.toml +++ b/vm/Cargo.toml @@ -70,6 +70,7 @@ num_enum = "0.4" smallbox = "0.8" bstr = "0.2.12" crossbeam-utils = "0.7" +generational-arena = "0.2" ## unicode stuff unicode_names2 = "0.4" diff --git a/vm/src/stdlib/random.rs b/vm/src/stdlib/random.rs index 63144451c..a057b224d 100644 --- a/vm/src/stdlib/random.rs +++ b/vm/src/stdlib/random.rs @@ -7,8 +7,9 @@ mod _random { use crate::function::OptionalOption; use crate::obj::objint::PyIntRef; use crate::obj::objtype::PyClassRef; - use crate::pyobject::{PyClassImpl, PyRef, PyResult, PyValue}; + use crate::pyobject::{PyClassImpl, PyRef, PyResult, PyValue, ThreadSafe}; use crate::VirtualMachine; + use generational_arena::{self, Arena}; use num_bigint::{BigInt, Sign}; use num_traits::Signed; use rand::RngCore; @@ -53,12 +54,51 @@ mod _random { } } + thread_local!(static RNG_HANDLES: RefCell> = RefCell::new(Arena::new())); + + #[derive(Debug)] + struct RngHandle(generational_arena::Index); + impl RngHandle { + fn new(rng: PyRng) -> Self { + let idx = RNG_HANDLES.with(|arena| arena.borrow_mut().insert(rng)); + RngHandle(idx) + } + fn exec(&self, func: F) -> R + where + F: Fn(&mut PyRng) -> R, + { + RNG_HANDLES.with(|arena| { + func( + arena + .borrow_mut() + .get_mut(self.0) + .expect("index was removed"), + ) + }) + } + fn replace(&self, rng: PyRng) { + RNG_HANDLES.with(|arena| { + *arena + .borrow_mut() + .get_mut(self.0) + .expect("index was removed") = rng + }) + } + } + impl Drop for RngHandle { + fn drop(&mut self) { + RNG_HANDLES.with(|arena| arena.borrow_mut().remove(self.0)); + } + } + #[pyclass(name = "Random")] #[derive(Debug)] struct PyRandom { - rng: RefCell, + rng: RngHandle, } + impl ThreadSafe for PyRandom {} + impl PyValue for PyRandom { fn class(vm: &VirtualMachine) -> PyClassRef { vm.class("_random", "Random") @@ -70,14 +110,14 @@ mod _random { #[pyslot(new)] fn new(cls: PyClassRef, vm: &VirtualMachine) -> PyResult> { PyRandom { - rng: RefCell::new(PyRng::default()), + rng: RngHandle::new(PyRng::default()), } .into_ref_with_type(vm, cls) } #[pymethod] fn random(&self) -> f64 { - mt19937::gen_res53(&mut *self.rng.borrow_mut()) + self.rng.exec(mt19937::gen_res53) } #[pymethod] @@ -93,31 +133,32 @@ mod _random { } }; - *self.rng.borrow_mut() = new_rng; + self.rng.replace(new_rng); } #[pymethod] - fn getrandbits(&self, mut k: usize) -> BigInt { - let mut rng = self.rng.borrow_mut(); + fn getrandbits(&self, k: usize) -> BigInt { + self.rng.exec(|rng| { + let mut k = k; + let mut gen_u32 = |k| rng.next_u32() >> (32 - k) as u32; - let mut gen_u32 = |k| rng.next_u32() >> (32 - k) as u32; + if k <= 32 { + return gen_u32(k).into(); + } - if k <= 32 { - return gen_u32(k).into(); - } + let words = (k - 1) / 8 + 1; + let mut wordarray = vec![0u32; words]; - 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; + } - 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) + BigInt::from_slice(Sign::NoSign, &wordarray) + }) } } }