Skip to content

Commit

Permalink
Let gen_range support ..len and ..=len for unsigned ints
Browse files Browse the repository at this point in the history
  • Loading branch information
dhardy committed Sep 6, 2024
1 parent 2144ae5 commit 7bdd833
Show file tree
Hide file tree
Showing 6 changed files with 50 additions and 14 deletions.
37 changes: 36 additions & 1 deletion src/distr/uniform.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -439,6 +439,41 @@ impl<T: SampleUniform + PartialOrd> SampleRange<T> for RangeInclusive<T> {
}
}

macro_rules! impl_sample_range_u {
($t:ty) => {
impl SampleRange<$t> for RangeTo<$t> {
#[inline]
fn sample_single<R: RngCore + ?Sized>(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<R: RngCore + ?Sized>(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::*;
Expand Down
9 changes: 5 additions & 4 deletions src/rng.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
///
Expand All @@ -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);
/// ```
///
Expand Down Expand Up @@ -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);
}
Expand Down
2 changes: 1 addition & 1 deletion src/seq/index.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
8 changes: 4 additions & 4 deletions src/seq/iterator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)),
};
}

Expand All @@ -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)
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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;
}
Expand Down
2 changes: 1 addition & 1 deletion src/seq/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
6 changes: 3 additions & 3 deletions src/seq/slice.rs
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ pub trait IndexedRandom: Index<usize> {
if self.is_empty() {
None
} else {
Some(&self[rng.gen_range(0..self.len())])
Some(&self[rng.gen_range(..self.len())])
}
}

Expand Down Expand Up @@ -259,7 +259,7 @@ pub trait IndexedMutRandom: IndexedRandom + IndexMut<usize> {
None
} else {
let len = self.len();
Some(&mut self[rng.gen_range(0..len)])
Some(&mut self[rng.gen_range(..len)])
}
}

Expand Down Expand Up @@ -410,7 +410,7 @@ impl<T> 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);
}
}
Expand Down

0 comments on commit 7bdd833

Please sign in to comment.