Jeff Vander Stoep | f00a459 | 2024-02-08 10:03:49 +0100 | [diff] [blame^] | 1 | //! A global, thread-local random number generator. |
| 2 | |
| 3 | use crate::Rng; |
| 4 | |
| 5 | use std::cell::Cell; |
| 6 | use std::ops::RangeBounds; |
| 7 | |
| 8 | // Chosen by fair roll of the dice. |
| 9 | const DEFAULT_RNG_SEED: u64 = 0xef6f79ed30ba75a; |
| 10 | |
| 11 | impl Default for Rng { |
| 12 | /// Initialize the `Rng` from the system's random number generator. |
| 13 | /// |
| 14 | /// This is equivalent to [`Rng::new()`]. |
| 15 | #[inline] |
| 16 | fn default() -> Rng { |
| 17 | Rng::new() |
| 18 | } |
| 19 | } |
| 20 | |
| 21 | impl Rng { |
| 22 | /// Creates a new random number generator. |
| 23 | #[inline] |
| 24 | pub fn new() -> Rng { |
| 25 | try_with_rng(Rng::fork).unwrap_or_else(|_| Rng::with_seed(0x4d595df4d0f33173)) |
| 26 | } |
| 27 | } |
| 28 | |
| 29 | thread_local! { |
| 30 | static RNG: Cell<Rng> = Cell::new(Rng(random_seed().unwrap_or(DEFAULT_RNG_SEED))); |
| 31 | } |
| 32 | |
| 33 | /// Run an operation with the current thread-local generator. |
| 34 | #[inline] |
| 35 | fn with_rng<R>(f: impl FnOnce(&mut Rng) -> R) -> R { |
| 36 | RNG.with(|rng| { |
| 37 | let current = rng.replace(Rng(0)); |
| 38 | |
| 39 | let mut restore = RestoreOnDrop { rng, current }; |
| 40 | |
| 41 | f(&mut restore.current) |
| 42 | }) |
| 43 | } |
| 44 | |
| 45 | /// Try to run an operation with the current thread-local generator. |
| 46 | #[inline] |
| 47 | fn try_with_rng<R>(f: impl FnOnce(&mut Rng) -> R) -> Result<R, std::thread::AccessError> { |
| 48 | RNG.try_with(|rng| { |
| 49 | let current = rng.replace(Rng(0)); |
| 50 | |
| 51 | let mut restore = RestoreOnDrop { rng, current }; |
| 52 | |
| 53 | f(&mut restore.current) |
| 54 | }) |
| 55 | } |
| 56 | |
| 57 | /// Make sure the original RNG is restored even on panic. |
| 58 | struct RestoreOnDrop<'a> { |
| 59 | rng: &'a Cell<Rng>, |
| 60 | current: Rng, |
| 61 | } |
| 62 | |
| 63 | impl Drop for RestoreOnDrop<'_> { |
| 64 | fn drop(&mut self) { |
| 65 | self.rng.set(Rng(self.current.0)); |
| 66 | } |
| 67 | } |
| 68 | |
| 69 | /// Initializes the thread-local generator with the given seed. |
| 70 | #[inline] |
| 71 | pub fn seed(seed: u64) { |
| 72 | with_rng(|r| r.seed(seed)); |
| 73 | } |
| 74 | |
| 75 | /// Gives back **current** seed that is being held by the thread-local generator. |
| 76 | #[inline] |
| 77 | pub fn get_seed() -> u64 { |
| 78 | with_rng(|r| r.get_seed()) |
| 79 | } |
| 80 | |
| 81 | /// Generates a random `bool`. |
| 82 | #[inline] |
| 83 | pub fn bool() -> bool { |
| 84 | with_rng(|r| r.bool()) |
| 85 | } |
| 86 | |
| 87 | /// Generates a random `char` in ranges a-z and A-Z. |
| 88 | #[inline] |
| 89 | pub fn alphabetic() -> char { |
| 90 | with_rng(|r| r.alphabetic()) |
| 91 | } |
| 92 | |
| 93 | /// Generates a random `char` in ranges a-z, A-Z and 0-9. |
| 94 | #[inline] |
| 95 | pub fn alphanumeric() -> char { |
| 96 | with_rng(|r| r.alphanumeric()) |
| 97 | } |
| 98 | |
| 99 | /// Generates a random `char` in range a-z. |
| 100 | #[inline] |
| 101 | pub fn lowercase() -> char { |
| 102 | with_rng(|r| r.lowercase()) |
| 103 | } |
| 104 | |
| 105 | /// Generates a random `char` in range A-Z. |
| 106 | #[inline] |
| 107 | pub fn uppercase() -> char { |
| 108 | with_rng(|r| r.uppercase()) |
| 109 | } |
| 110 | |
| 111 | /// Choose an item from an iterator at random. |
| 112 | /// |
| 113 | /// This function may have an unexpected result if the `len()` property of the |
| 114 | /// iterator does not match the actual number of items in the iterator. If |
| 115 | /// the iterator is empty, this returns `None`. |
| 116 | #[inline] |
| 117 | pub fn choice<I>(iter: I) -> Option<I::Item> |
| 118 | where |
| 119 | I: IntoIterator, |
| 120 | I::IntoIter: ExactSizeIterator, |
| 121 | { |
| 122 | with_rng(|r| r.choice(iter)) |
| 123 | } |
| 124 | |
| 125 | /// Generates a random digit in the given `base`. |
| 126 | /// |
| 127 | /// Digits are represented by `char`s in ranges 0-9 and a-z. |
| 128 | /// |
| 129 | /// Panics if the base is zero or greater than 36. |
| 130 | #[inline] |
| 131 | pub fn digit(base: u32) -> char { |
| 132 | with_rng(|r| r.digit(base)) |
| 133 | } |
| 134 | |
| 135 | /// Shuffles a slice randomly. |
| 136 | #[inline] |
| 137 | pub fn shuffle<T>(slice: &mut [T]) { |
| 138 | with_rng(|r| r.shuffle(slice)) |
| 139 | } |
| 140 | |
| 141 | macro_rules! integer { |
| 142 | ($t:tt, $doc:tt) => { |
| 143 | #[doc = $doc] |
| 144 | /// |
| 145 | /// Panics if the range is empty. |
| 146 | #[inline] |
| 147 | pub fn $t(range: impl RangeBounds<$t>) -> $t { |
| 148 | with_rng(|r| r.$t(range)) |
| 149 | } |
| 150 | }; |
| 151 | } |
| 152 | |
| 153 | integer!(u8, "Generates a random `u8` in the given range."); |
| 154 | integer!(i8, "Generates a random `i8` in the given range."); |
| 155 | integer!(u16, "Generates a random `u16` in the given range."); |
| 156 | integer!(i16, "Generates a random `i16` in the given range."); |
| 157 | integer!(u32, "Generates a random `u32` in the given range."); |
| 158 | integer!(i32, "Generates a random `i32` in the given range."); |
| 159 | integer!(u64, "Generates a random `u64` in the given range."); |
| 160 | integer!(i64, "Generates a random `i64` in the given range."); |
| 161 | integer!(u128, "Generates a random `u128` in the given range."); |
| 162 | integer!(i128, "Generates a random `i128` in the given range."); |
| 163 | integer!(usize, "Generates a random `usize` in the given range."); |
| 164 | integer!(isize, "Generates a random `isize` in the given range."); |
| 165 | integer!(char, "Generates a random `char` in the given range."); |
| 166 | |
| 167 | /// Generates a random `f32` in range `0..1`. |
| 168 | pub fn f32() -> f32 { |
| 169 | with_rng(|r| r.f32()) |
| 170 | } |
| 171 | |
| 172 | /// Generates a random `f64` in range `0..1`. |
| 173 | pub fn f64() -> f64 { |
| 174 | with_rng(|r| r.f64()) |
| 175 | } |
| 176 | |
| 177 | /// Collects `amount` values at random from the iterator into a vector. |
| 178 | pub fn choose_multiple<T: Iterator>(source: T, amount: usize) -> Vec<T::Item> { |
| 179 | with_rng(|rng| rng.choose_multiple(source, amount)) |
| 180 | } |
| 181 | |
| 182 | #[cfg(not(all( |
| 183 | any(target_arch = "wasm32", target_arch = "wasm64"), |
| 184 | target_os = "unknown" |
| 185 | )))] |
| 186 | fn random_seed() -> Option<u64> { |
| 187 | use std::collections::hash_map::DefaultHasher; |
| 188 | use std::hash::{Hash, Hasher}; |
| 189 | use std::thread; |
| 190 | use std::time::Instant; |
| 191 | |
| 192 | let mut hasher = DefaultHasher::new(); |
| 193 | Instant::now().hash(&mut hasher); |
| 194 | thread::current().id().hash(&mut hasher); |
| 195 | let hash = hasher.finish(); |
| 196 | Some((hash << 1) | 1) |
| 197 | } |
| 198 | |
| 199 | #[cfg(all( |
| 200 | any(target_arch = "wasm32", target_arch = "wasm64"), |
| 201 | target_os = "unknown", |
| 202 | feature = "js" |
| 203 | ))] |
| 204 | fn random_seed() -> Option<u64> { |
| 205 | // TODO(notgull): Failures should be logged somewhere. |
| 206 | let mut seed = [0u8; 8]; |
| 207 | getrandom::getrandom(&mut seed).ok()?; |
| 208 | Some(u64::from_ne_bytes(seed)) |
| 209 | } |
| 210 | |
| 211 | #[cfg(all( |
| 212 | any(target_arch = "wasm32", target_arch = "wasm64"), |
| 213 | target_os = "unknown", |
| 214 | not(feature = "js") |
| 215 | ))] |
| 216 | fn random_seed() -> Option<u64> { |
| 217 | None |
| 218 | } |