Skip to content

Commit

Permalink
Implement Distribution for Range(Inclusive)
Browse files Browse the repository at this point in the history
  • Loading branch information
Ralith committed Jul 1, 2019
1 parent 9f119ff commit b25b657
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 4 deletions.
31 changes: 29 additions & 2 deletions src/distributions/uniform.rs
Expand Up @@ -111,6 +111,7 @@
use std::time::Duration;
#[cfg(not(feature = "std"))]
use core::time::Duration;
use core::ops::{Range, RangeInclusive};

use crate::Rng;
use crate::distributions::Distribution;
Expand Down Expand Up @@ -147,6 +148,9 @@ use packed_simd::*;
/// `Uniform::new(low, high)`, i.e., excluding `high`. In particular care must
/// be taken to ensure that rounding never results values `< low` or `>= high`.
///
/// Range expressions like `0..10` can be used as a `Uniform` distribution,
/// but are less efficient if multiple samples are taken.
///
/// # Example
///
/// ```
Expand Down Expand Up @@ -278,18 +282,30 @@ pub trait UniformSampler: Sized {
}
}

impl<X: SampleUniform> From<::core::ops::Range<X>> for Uniform<X> {
impl<X: SampleUniform> From<Range<X>> for Uniform<X> {
fn from(r: ::core::ops::Range<X>) -> Uniform<X> {
Uniform::new(r.start, r.end)
}
}

impl<X: SampleUniform> From<::core::ops::RangeInclusive<X>> for Uniform<X> {
impl<X: SampleUniform> From<RangeInclusive<X>> for Uniform<X> {
fn from(r: ::core::ops::RangeInclusive<X>) -> Uniform<X> {
Uniform::new_inclusive(r.start(), r.end())
}
}

impl<T: SampleUniform> Distribution<T> for Range<T> {
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> T {
T::Sampler::sample_single(&self.start, &self.end, rng)
}
}

impl<T: SampleUniform> Distribution<T> for RangeInclusive<T> {
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> T {
T::Sampler::sample_single_inclusive(self.start(), self.end(), rng)
}
}

/// Helper trait similar to [`Borrow`] but implemented
/// only for SampleUniform and references to SampleUniform in
/// order to resolve ambiguity issues.
Expand Down Expand Up @@ -1274,6 +1290,17 @@ mod tests {
assert_eq!(r.inner.scale, 5.0);
}

#[test]
fn test_std_range_distribution() {
let mut rng = crate::test::rng(474);
for _ in 0..100 {
let x = rng.sample(0..10);
assert!(x >= 0 && x < 10);
let x = rng.sample(0..=10);
assert!(x >= 0 && x <= 10);
}
}

#[test]
fn test_uniform_from_std_range_inclusive() {
let r = Uniform::from(2u32..=6);
Expand Down
5 changes: 3 additions & 2 deletions src/lib.rs
Expand Up @@ -227,9 +227,10 @@ pub trait Rng: RngCore {
/// use rand::distributions::Uniform;
///
/// let mut rng = thread_rng();
/// let x = rng.sample(Uniform::new(10u32, 15));
/// let x = rng.sample(10u32..15);
/// // Type annotation requires two types, the type and distribution; the
/// // distribution can be inferred.
/// // distribution can be inferred. `Uniform` is more efficient than the
/// // above range syntax if multiple samples are taken.
/// let y = rng.sample::<u16, _>(Uniform::new(10, 15));
/// ```
fn sample<T, D: Distribution<T>>(&mut self, distr: D) -> T {
Expand Down

0 comments on commit b25b657

Please sign in to comment.