Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Rng::iter #275

Closed
wants to merge 3 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 54 additions & 1 deletion benches/misc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ extern crate rand;

use test::{black_box, Bencher};

use rand::{Rng, weak_rng};
use rand::{Rng, RngCore, weak_rng};
use rand::seq::*;

#[bench]
Expand Down Expand Up @@ -60,3 +60,56 @@ macro_rules! sample_indices {
sample_indices!(misc_sample_indices_10_of_1k, 10, 1000);
sample_indices!(misc_sample_indices_50_of_1k, 50, 1000);
sample_indices!(misc_sample_indices_100_of_1k, 100, 1000);

#[bench]
fn gen_1k_iter_repeat(b: &mut Bencher) {
use std::iter;
let mut rng = weak_rng();
b.iter(|| {
let v: Vec<u32> = iter::repeat(()).map(|()| rng.next_u32()).take(256).collect();
black_box(v);
});
b.bytes = 1024;
}

#[allow(deprecated)]
#[bench]
fn gen_1k_gen_iter(b: &mut Bencher) {
let mut rng = weak_rng();
b.iter(|| {
let v: Vec<u32> = rng.gen_iter().take(256).collect();
black_box(v);
});
b.bytes = 1024;
}

#[bench]
fn gen_1k_iter1(b: &mut Bencher) {
let mut rng = weak_rng();
b.iter(|| {
let v: Vec<u32> = rng.iter().take(256).map(|rng| rng.next_u32()).collect();
black_box(v);
});
b.bytes = 1024;
}

#[bench]
fn gen_1k_iter2(b: &mut Bencher) {
let mut rng = weak_rng();
b.iter(|| {
let v: Vec<u32> = rng.iter().map(|rng| rng.next_u32()).take(256).collect();
black_box(v);
});
b.bytes = 1024;
}

#[bench]
fn gen_1k_fill(b: &mut Bencher) {
let mut rng = weak_rng();
let mut buf = [0u32; 256];
b.iter(|| {
rng.fill(&mut buf[..]);
black_box(buf);
});
b.bytes = 1024;
}
227 changes: 227 additions & 0 deletions src/iter.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,227 @@
// Copyright 2017-2018 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// https://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Iterators over `RngCore`

use core::cmp::min;
use core::usize;

use RngCore;

// TODO: in the future (see https://github.com/rust-lang/rfcs/issues/1403)
// it may be possible to allow reborrows in user code; this would let us
// replace `rng: &'a mut R` with `rng: R` in `Iter`, without having to create a
// redundant reference when calling `Map::next`. In this case `Rng::iter` would
// return `Iter<&mut Self>` and a separate constructor would be needed for
// reference/copy types not needing an extra reference.

/// An "iterator" over a random number generator; created by [`Rng::iter`].
///
/// This does not implement `std::iter::Iterator` since we cannot support
/// `next()`: it makes no sense to return a copy of an RNG, and though in
/// theory it should be possible to return `&mut RngCore`, `Iterator` does not
/// allow the lifetime of the item returned by `next` to be linked to the
/// iterator (or enclosed RNG). Instead, we support other selected operations
/// such as `map` and `take` directly.
///
/// [`Rng::iter`]: ../trait.Rng.html#method.iter
#[derive(Debug)]
pub struct Iter<'a, R: RngCore + ?Sized + 'a> {
pub(crate) rng: &'a mut R,
pub(crate) len: Option<usize>,
}

impl<'a, R: RngCore + ?Sized + 'a> Iter<'a, R> {
pub(crate) fn new(rng: &'a mut R) -> Iter<'a, R> {
Iter { rng, len: None }
}
}

impl<'a, R: RngCore + ?Sized + 'a> Iter<'a, R> {
/// Restricts the number of generated items to at most `len`.
pub fn take(self, len: usize) -> Self {
Iter {
rng: self.rng,
len: Some(self.len.map_or(len, |old| min(old, len))),
}
}

/// Takes a closure and creates an iterator which calls that closure on
/// each element.
///
/// ### Example
///
/// ```rust
/// use rand::{thread_rng, Rng};
/// use rand::distributions::Range;
///
/// let die_range = Range::new(1, 7);
/// let mut rng = thread_rng();
/// let mut die = rng.iter().map(|rng| rng.sample(die_range));
/// for _ in 0..3 {
/// println!("Die roll: {}", die.next().unwrap());
/// }
/// ```
pub fn map<B, F>(self, f: F) -> Map<'a, R, B, F>
where F: FnMut(&mut R) -> B
{
Map {
rng: self.rng,
len: self.len,
f: f,
}
}

/// Creates an iterator that works like map, but flattens nested structure.
///
/// The [`map`] adapter is very useful, but only when the closure argument
/// produces values. If it produces an iterator instead, there's an extra
/// layer of indirection. `flat_map()` will remove this extra layer on its
/// own.
///
/// ### Example
///
/// ```rust
/// use rand::{thread_rng, Rng};
/// use rand::distributions::Range;
///
/// let len_range = Range::new(1, 10);
/// let mut rng = thread_rng();
///
/// // Count from 1 to a number between 1 and 9 several times:
/// let mut iter = rng.iter().flat_map(|rng| 1..rng.sample(len_range)).take(20);
/// while let Some(n) = iter.next() {
/// println!("{}", n);
/// }
/// ```
///
/// [`map`]: struct.Iter.html#method.map
pub fn flat_map<U, F>(self, f: F) -> FlatMap<'a, R, U, F>
where F: FnMut(&mut R) -> U, U: IntoIterator
{
FlatMap {
rng: self.rng,
len: self.len,
f: f,
frontiter: None,
}
}
}

/// Type created by [`Iter::map`](struct.Iter.html#method.map)
#[derive(Debug)]
pub struct Map<'a, R:?Sized+'a, B, F> where F: FnMut(&mut R) -> B {
rng: &'a mut R,
len: Option<usize>,
f: F,
}
impl<'a, R:?Sized+'a, B, F> Iterator for Map<'a, R, B, F>
where F: FnMut(&mut R) -> B
{
type Item = B;

fn next(&mut self) -> Option<B> {
match self.len {
Some(0) => return None,
Some(ref mut n) => { *n -= 1; }
None => {}
}

Some((self.f)(self.rng))
}

fn size_hint(&self) -> (usize, Option<usize>) {
// If len == None we have an infinite iterator; usize::MAX is nearest
// available lower bound. Probably this suffices to make the following equal:
// rng.iter().take(n).map(f).size_hint() == rng.iter().map(f).take(n).size_hint()
self.len.map_or((usize::MAX, None), |len| (len, Some(len)))
}
}

/// Type created by [`Iter::flat_map`](struct.Iter.html#method.flat_map)
#[derive(Debug)]
pub struct FlatMap<'a, R:?Sized+'a, U, F>
where F: FnMut(&mut R) -> U, U: IntoIterator
{
rng: &'a mut R,
len: Option<usize>,
f: F,
frontiter: Option<U::IntoIter>,
}
impl<'a, R:?Sized+'a, U, F> Iterator for FlatMap<'a, R, U, F>
where F: FnMut(&mut R) -> U, U: IntoIterator
{
type Item = <U as IntoIterator>::Item;

fn next(&mut self) -> Option<Self::Item> {
loop {
if let Some(ref mut inner) = self.frontiter {
if let Some(x) = inner.by_ref().next() {
return Some(x)
}
}

match self.len {
Some(0) => return None,
Some(ref mut n) => { *n -= 1; }
None => {}
}

self.frontiter = Some(IntoIterator::into_iter((self.f)(self.rng)));
}
}

fn size_hint(&self) -> (usize, Option<usize>) {
if self.len == Some(0) {
// No new iters, so we have frontiter or nothing
self.frontiter.as_ref().map_or((0, Some(0)), |it| it.size_hint())
} else {
// Can't compute an actual bound without producing the sub-iters,
// which we don't want to do. But we may have a lower bound.
let lb = self.frontiter.as_ref().map_or(0, |it| it.size_hint().0);
(lb, None)
}
}
}

#[cfg(test)]
mod tests {
use {Rng, RngCore};
use distributions::{Uniform};
#[cfg(all(not(feature="std"), feature="alloc"))] use alloc::{Vec, String};

#[test]
#[cfg(any(feature="std", feature="alloc"))]
fn test_iter() {
let mut rng = ::test::rng(160);

let x: Vec<()> = rng.iter().take(10).map(|_| ()).collect();
assert_eq!(x.len(), 10);
let y: Vec<u32> = rng.iter().take(10).map(|rng| rng.sample(Uniform)).collect();
assert_eq!(y.len(), 10);
let z: Vec<u32> = rng.iter().take(10).flat_map(|rng|
vec![rng.sample(Uniform), rng.sample(Uniform)].into_iter()).collect();
assert_eq!(z.len(), 20);
let w: Vec<String> = rng.iter().take(10).flat_map(|_| vec![].into_iter()).collect();
assert_eq!(w.len(), 0);
}

#[test]
fn test_dyn_dispatch() {
let mut rng = ::test::rng(161);
let mut r: &mut RngCore = &mut rng;

let mut x = 0;
for n in r.iter().map(|rng| rng.next_u32()).take(2) {
x ^= n;
}
assert!(x != 0);
}
}
45 changes: 40 additions & 5 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -253,7 +253,7 @@
#![cfg_attr(all(target_arch = "wasm32", not(target_os = "emscripten")), recursion_limit="128")]

#[cfg(feature="std")] extern crate std as core;
#[cfg(all(feature = "alloc", not(feature="std")))] extern crate alloc;
#[cfg(all(feature = "alloc", not(feature="std")))] #[macro_use] extern crate alloc;

#[cfg(test)] #[cfg(feature="serde-1")] extern crate bincode;
#[cfg(feature="serde-1")] extern crate serde;
Expand Down Expand Up @@ -297,6 +297,7 @@ use distributions::range::SampleRange;
// public modules
pub mod distributions;
mod impls;
pub mod iter;
pub mod jitter;
pub mod mock;
#[cfg(feature="std")] pub mod os;
Expand All @@ -315,9 +316,9 @@ pub mod isaac {
}

// private modules
mod le;
#[cfg(feature="std")] mod entropy_rng;
mod error;
mod le;
mod prng;
#[cfg(feature="std")] mod thread_rng;

Expand Down Expand Up @@ -601,6 +602,8 @@ pub trait Rng: RngCore + Sized {
/// println!("{:?}", rng.gen_iter::<(f64, bool)>().take(5)
/// .collect::<Vec<(f64, bool)>>());
/// ```
#[allow(deprecated)]
#[deprecated(since="0.5.0", note="replaced by Rng::iter")]
fn gen_iter<T>(&mut self) -> Generator<T, &mut Self> where Uniform: Distribution<T> {
Generator { rng: self, _marker: marker::PhantomData }
}
Expand Down Expand Up @@ -654,6 +657,35 @@ pub trait Rng: RngCore + Sized {
n <= 1 || self.gen_range(0, n) == 0
}

/// Construct an iterator on an `Rng`.
///
/// ### Example
///
/*
/// ```rust
/// use rand::{thread_rng, Rng};
/// use distributions::Range;
///
/// let die_range = Range::new(1, 7);
/// let mut die = thread_rng().iter().map(|rng| rng.sample(die_range));
/// for _ in 0..3 {
/// println!("Die roll: {}", die.next());
/// }
/// ```
*/
/* TODO: Alphanumeric
/// ```rust
/// use rand::{thread_rng, Rng};
/// use rand::distributions::Alphanumeric;
///
/// let mut rng = thread_rng();
/// let x: String = rng.iter().map(|rng| rng.sample(Alphanumeric)).take(6).collect();
/// ```
*/
fn iter<'a>(&'a mut self) -> iter::Iter<'a, Self> {
iter::Iter::new(self)
}

/// Return an iterator of random characters from the set A-Z,a-z,0-9.
///
/// # Example
Expand Down Expand Up @@ -861,11 +893,14 @@ impl_as_byte_slice_arrays!(32, N,N,N,N,N,N,N,N,N,N,N,N,N,N,N,N,N,N,N,N,N,N,N,N,N
/// [`gen_iter`]: trait.Rng.html#method.gen_iter
/// [`Rng`]: trait.Rng.html
#[derive(Debug)]
#[allow(deprecated)]
#[deprecated(since="0.5.0", note="replaced by Rng::iter")]
pub struct Generator<T, R: RngCore> {
rng: R,
_marker: marker::PhantomData<fn() -> T>,
}

#[allow(deprecated)]
impl<T, R: RngCore> Iterator for Generator<T, R> where Uniform: Distribution<T> {
type Item = T;

Expand Down Expand Up @@ -1222,9 +1257,9 @@ mod test {
#[test]
fn test_gen_vec() {
let mut r = rng(106);
assert_eq!(r.gen_iter::<u8>().take(0).count(), 0);
assert_eq!(r.gen_iter::<u8>().take(10).count(), 10);
assert_eq!(r.gen_iter::<f64>().take(16).count(), 16);
assert_eq!(r.iter().map(|rng| rng.gen::<u8>()).take(0).count(), 0);
assert_eq!(r.iter().map(|rng| rng.gen::<u8>()).take(10).count(), 10);
assert_eq!(r.iter().map(|rng| rng.gen::<f64>()).take(16).count(), 16);
}

#[test]
Expand Down