From c8745e7e78f7eeafd3d303a2febf42cf64917dd1 Mon Sep 17 00:00:00 2001 From: Diggory Hardy Date: Mon, 26 Aug 2024 09:42:59 +0100 Subject: [PATCH] Let gen_range support ..len and ..=len for unsigned ints --- src/distr/uniform.rs | 37 ++++++++++++++++++++++++++++++++++++- src/rng.rs | 9 +++++---- src/seq/index.rs | 2 +- src/seq/iterator.rs | 8 ++++---- src/seq/mod.rs | 2 +- src/seq/slice.rs | 6 +++--- 6 files changed, 50 insertions(+), 14 deletions(-) diff --git a/src/distr/uniform.rs b/src/distr/uniform.rs index 402efa7bd4..86a08fdc59 100644 --- a/src/distr/uniform.rs +++ b/src/distr/uniform.rs @@ -120,7 +120,7 @@ mod other; pub use other::{UniformChar, UniformDuration}; use core::fmt; -use core::ops::{Range, RangeInclusive}; +use core::ops::{Range, RangeInclusive, RangeTo, RangeToInclusive}; use crate::distr::Distribution; use crate::{Rng, RngCore}; @@ -439,6 +439,41 @@ impl SampleRange for RangeInclusive { } } +macro_rules! impl_sample_range_u { + ($t:ty) => { + impl SampleRange<$t> for RangeTo<$t> { + #[inline] + fn sample_single(self, rng: &mut R) -> Result<$t, Error> { + <$t as SampleUniform>::Sampler::sample_single(0, self.end, rng) + } + + #[inline] + fn is_empty(&self) -> bool { + 0 == self.end + } + } + + impl SampleRange<$t> for RangeToInclusive<$t> { + #[inline] + fn sample_single(self, rng: &mut R) -> Result<$t, Error> { + <$t as SampleUniform>::Sampler::sample_single_inclusive(0, self.end, rng) + } + + #[inline] + fn is_empty(&self) -> bool { + false + } + } + }; +} + +impl_sample_range_u!(u8); +impl_sample_range_u!(u16); +impl_sample_range_u!(u32); +impl_sample_range_u!(u64); +impl_sample_range_u!(u128); +impl_sample_range_u!(usize); + #[cfg(test)] mod tests { use super::*; diff --git a/src/rng.rs b/src/rng.rs index 913ca749b7..9190747a29 100644 --- a/src/rng.rs +++ b/src/rng.rs @@ -102,7 +102,8 @@ pub trait Rng: RngCore { /// made from the given range. See also the [`Uniform`] distribution /// type which may be faster if sampling from the same range repeatedly. /// - /// Only `gen_range(low..high)` and `gen_range(low..=high)` are supported. + /// All types support `low..high_exclusive` and `low..=high` range syntax. + /// Unsigned integer types also support `..high_exclusive` and `..=high` syntax. /// /// # Panics /// @@ -116,13 +117,13 @@ pub trait Rng: RngCore { /// let mut rng = thread_rng(); /// /// // Exclusive range - /// let n: u32 = rng.gen_range(0..10); + /// let n: u32 = rng.gen_range(..10); /// println!("{}", n); /// let m: f64 = rng.gen_range(-40.0..1.3e5); /// println!("{}", m); /// /// // Inclusive range - /// let n: u32 = rng.gen_range(0..=10); + /// let n: u32 = rng.gen_range(..=10); /// println!("{}", n); /// ``` /// @@ -500,7 +501,7 @@ mod test { let a: u32 = r.gen_range(12..=24); assert!((12..=24).contains(&a)); - assert_eq!(r.gen_range(0u32..1), 0u32); + assert_eq!(r.gen_range(..1u32), 0u32); assert_eq!(r.gen_range(-12i64..-11), -12i64); assert_eq!(r.gen_range(3_000_000..3_000_001), 3_000_000); } diff --git a/src/seq/index.rs b/src/seq/index.rs index 4932617f97..5bb1a7597f 100644 --- a/src/seq/index.rs +++ b/src/seq/index.rs @@ -434,7 +434,7 @@ where debug_assert!(amount <= length); let mut indices = Vec::with_capacity(amount as usize); for j in length - amount..length { - let t = rng.gen_range(0..=j); + let t = rng.gen_range(..=j); if let Some(pos) = indices.iter().position(|&x| x == t) { indices[pos] = j; } diff --git a/src/seq/iterator.rs b/src/seq/iterator.rs index ac25bb806e..148093157d 100644 --- a/src/seq/iterator.rs +++ b/src/seq/iterator.rs @@ -70,7 +70,7 @@ pub trait IteratorRandom: Iterator + Sized { return match lower { 0 => None, 1 => self.next(), - _ => self.nth(rng.gen_range(0..lower)), + _ => self.nth(rng.gen_range(..lower)), }; } @@ -80,7 +80,7 @@ pub trait IteratorRandom: Iterator + Sized { // Continue until the iterator is exhausted loop { if lower > 1 { - let ix = coin_flipper.rng.gen_range(0..lower + consumed); + let ix = coin_flipper.rng.gen_range(..lower + consumed); let skip = if ix < lower { result = self.nth(ix); lower - (ix + 1) @@ -203,7 +203,7 @@ pub trait IteratorRandom: Iterator + Sized { // Continue, since the iterator was not exhausted for (i, elem) in self.enumerate() { - let k = rng.gen_range(0..i + 1 + amount); + let k = rng.gen_range(..i + 1 + amount); if let Some(slot) = buf.get_mut(k) { *slot = elem; } @@ -239,7 +239,7 @@ pub trait IteratorRandom: Iterator + Sized { // If the iterator stops once, then so do we. if reservoir.len() == amount { for (i, elem) in self.enumerate() { - let k = rng.gen_range(0..i + 1 + amount); + let k = rng.gen_range(..i + 1 + amount); if let Some(slot) = reservoir.get_mut(k) { *slot = elem; } diff --git a/src/seq/mod.rs b/src/seq/mod.rs index e268c10c99..66459fe174 100644 --- a/src/seq/mod.rs +++ b/src/seq/mod.rs @@ -69,7 +69,7 @@ pub mod index { // Floyd's algorithm let mut indices = [0; N]; for (i, j) in (len - N..len).enumerate() { - let t = rng.gen_range(0..j + 1); + let t = rng.gen_range(..j + 1); if let Some(pos) = indices[0..i].iter().position(|&x| x == t) { indices[pos] = j; } diff --git a/src/seq/slice.rs b/src/seq/slice.rs index 72e1f8d1c0..930e545094 100644 --- a/src/seq/slice.rs +++ b/src/seq/slice.rs @@ -57,7 +57,7 @@ pub trait IndexedRandom: Index { if self.is_empty() { None } else { - Some(&self[rng.gen_range(0..self.len())]) + Some(&self[rng.gen_range(..self.len())]) } } @@ -259,7 +259,7 @@ pub trait IndexedMutRandom: IndexedRandom + IndexMut { None } else { let len = self.len(); - Some(&mut self[rng.gen_range(0..len)]) + Some(&mut self[rng.gen_range(..len)]) } } @@ -410,7 +410,7 @@ impl SliceRandom for [T] { } } else { for i in m..self.len() { - let index = rng.gen_range(0..i + 1); + let index = rng.gen_range(..i + 1); self.swap(i, index); } }