Skip to content

Commit

Permalink
Accept any IntoIterator in choose_multiple
Browse files Browse the repository at this point in the history
  • Loading branch information
tranzystorekk authored and notgull committed Nov 17, 2024
1 parent 1b93479 commit a69b481
Show file tree
Hide file tree
Showing 3 changed files with 11 additions and 10 deletions.
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ let elem = v[i];
Sample values from an array with `O(n)` complexity (`n` is the length of array):

```rust
fastrand::choose_multiple(vec![1, 4, 5].iter(), 2);
fastrand::choose_multiple([1, 4, 5], 2);
fastrand::choose_multiple(0..20, 12);
```

Expand Down
4 changes: 2 additions & 2 deletions src/global_rng.rs
Original file line number Diff line number Diff line change
Expand Up @@ -181,8 +181,8 @@ pub fn f64() -> f64 {
with_rng(|r| r.f64())
}

/// Collects `amount` values at random from the iterator into a vector.
pub fn choose_multiple<T: Iterator>(source: T, amount: usize) -> Vec<T::Item> {
/// Collects `amount` values at random from the iterable into a vector.
pub fn choose_multiple<I: IntoIterator>(source: I, amount: usize) -> Vec<I::Item> {
with_rng(|rng| rng.choose_multiple(source, amount))
}

Expand Down
15 changes: 8 additions & 7 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@
//! Sample values from an array with `O(n)` complexity (`n` is the length of array):
//!
//! ```
//! fastrand::choose_multiple(vec![1, 4, 5].iter(), 2);
//! fastrand::choose_multiple([1, 4, 5], 2);
//! fastrand::choose_multiple(0..20, 12);
//! ```
//!
Expand Down Expand Up @@ -378,27 +378,28 @@ impl Rng {
f64::from_bits((1 << (b - 2)) - (1 << f) + (self.u64(..) >> (b - f))) - 1.0
}

/// Collects `amount` values at random from the iterator into a vector.
/// Collects `amount` values at random from the iterable into a vector.
///
/// The length of the returned vector equals `amount` unless the iterator
/// The length of the returned vector equals `amount` unless the iterable
/// contains insufficient elements, in which case it equals the number of
/// elements available.
///
/// Complexity is `O(n)` where `n` is the length of the iterator.
/// Complexity is `O(n)` where `n` is the length of the iterable.
#[cfg(feature = "alloc")]
#[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
pub fn choose_multiple<T: Iterator>(&mut self, mut source: T, amount: usize) -> Vec<T::Item> {
pub fn choose_multiple<I: IntoIterator>(&mut self, source: I, amount: usize) -> Vec<I::Item> {
// Adapted from: https://docs.rs/rand/latest/rand/seq/trait.IteratorRandom.html#method.choose_multiple
let mut reservoir = Vec::with_capacity(amount);
let mut iter = source.into_iter();

reservoir.extend(source.by_ref().take(amount));
reservoir.extend(iter.by_ref().take(amount));

// Continue unless the iterator was exhausted
//
// note: this prevents iterators that "restart" from causing problems.
// If the iterator stops once, then so do we.
if reservoir.len() == amount {
for (i, elem) in source.enumerate() {
for (i, elem) in iter.enumerate() {
let end = i + 1 + amount;
let k = self.usize(0..end);
if let Some(slot) = reservoir.get_mut(k) {
Expand Down

0 comments on commit a69b481

Please sign in to comment.