simplify code

Author: newpavlov

thread-rng/src/lib.rs

@@ -13,10 +13,39 @@ use std::fmt;
 use std::rc::Rc;
 use std::thread_local;
 
-pub use rand_core::{self, CryptoRng, RngCore};
+pub use rand_core;
 
-mod reseeding_rng;
-use reseeding_rng::ReseedingRng;
+use rand_chacha::ChaCha12Rng;
+use rand_core::{CryptoRng, RngCore, SeedableRng};
+
+// Number of generated bytes after which to reseed `ThreadRng`.
+// According to benchmarks, reseeding has a noticeable impact with thresholds
+// of 32 kB and less. We choose 64 kB to avoid significant overhead.
+const RESEED_THRESHOLD: isize = 1024 * 64;
+
+struct InnerState {
+    rng: ChaCha12Rng,
+    bytes_until_reseed: isize,
+}
+
+impl InnerState {
+    #[inline(always)]
+    fn reseed(&mut self) -> Result<(), rand_core::getrandom::Error> {
+        self.bytes_until_reseed = RESEED_THRESHOLD;
+        self.rng = ChaCha12Rng::try_from_os_rng()?;
+        Ok(())
+    }
+
+    #[inline(always)]
+    fn reseed_check(&mut self, n: isize) {
+        if self.bytes_until_reseed < 0 {
+            // If system RNG has failed for some reason, ignore the error
+            // and continue to work with the old RNG state.
+            let _ = self.reseed();
+        }
+        self.bytes_until_reseed -= n;
+    }
+}
 
 thread_local!(
     // We require Rc<..> to avoid premature freeing when ThreadRng is used
@@ -35,11 +64,12 @@ thread_local!(
     // completely under our control. We just have to ensure none of them use
     // `ThreadRng` internally, which is nonsensical anyway. We should also never run
     // `ThreadRng` in destructors of its implementation, which is also nonsensical.
-    static THREAD_RNG_KEY: Rc<UnsafeCell<ReseedingRng>> = {
-        match ReseedingRng::new() {
-            Ok(rng) => Rc::new(UnsafeCell::new(rng)),
-            Err(err) => panic!("could not initialize ThreadRng: {}", err),
-        }
+    static THREAD_RNG_KEY: Rc<UnsafeCell<InnerState>> = {
+        let rng = match ChaCha12Rng::try_from_os_rng() {
+            Ok(rng) => rng,
+            Err(err) => panic!("could not initialize ThreadRng: {err}"),
+        };
+        Rc::new(UnsafeCell::new(InnerState { rng, bytes_until_reseed: RESEED_THRESHOLD }))
     }
 );
 
@@ -95,24 +125,17 @@ thread_local!(
 #[derive(Clone)]
 pub struct ThreadRng {
     // Rc is explicitly !Send and !Sync
-    rng: Rc<UnsafeCell<ReseedingRng>>,
+    rng: Rc<UnsafeCell<InnerState>>,
 }
 
 impl ThreadRng {
     /// Immediately reseed the generator
     ///
     /// This discards any remaining random data in the cache.
     pub fn reseed(&mut self) -> Result<(), rand_core::getrandom::Error> {
-        // SAFETY: We must make sure to stop using `rng` before anyone else
-        // creates another mutable reference
-        let rng = unsafe { &mut *self.rng.get() };
-        rng.reseed()
-    }
-
-    fn get_rng(&mut self) -> &mut ReseedingRng {
-        // SAFETY: We must make sure to stop using `rng` before anyone else
-        // creates another mutable reference
-        unsafe { &mut *self.rng.get() }
+        // SAFETY: The state is thread-local and the reference does not leak from this method
+        let s = unsafe { &mut *self.rng.get() };
+        s.reseed()
     }
 }
 
@@ -158,17 +181,28 @@ impl Default for ThreadRng {
 impl RngCore for ThreadRng {
     #[inline(always)]
     fn next_u32(&mut self) -> u32 {
-        self.get_rng().next_u32()
+        // SAFETY: The state is thread-local and the reference does not leak from this method
+        let s = unsafe { &mut *self.rng.get() };
+        s.reseed_check(core::mem::size_of::<u32>() as isize);
+        s.rng.next_u32()
     }
 
     #[inline(always)]
     fn next_u64(&mut self) -> u64 {
-        self.get_rng().next_u64()
+        // SAFETY: The state is thread-local and the reference does not leak from this method
+        let s = unsafe { &mut *self.rng.get() };
+        s.reseed_check(core::mem::size_of::<u64>() as isize);
+        s.rng.next_u64()
     }
 
     #[inline(always)]
     fn fill_bytes(&mut self, dest: &mut [u8]) {
-        self.get_rng().fill_bytes(dest)
+        // SAFETY: The state is thread-local and the reference does not leak from this method
+        let s = unsafe { &mut *self.rng.get() };
+        // Valid allocations can not be bigger than `isize::MAX` bytes,
+        // so we can cast length to `isize` without issues.
+        s.reseed_check(dest.len() as isize);
+        s.rng.fill_bytes(dest)
     }
 }