Add choose_multple function for choosing several values (#55)

Author: notgull

.github/workflows/ci.yml

@@ -52,6 +52,9 @@ jobs:
       - name: Build with no default features
         # Use no-std target to ensure we don't link to std.
         run: cargo build --no-default-features --target thumbv7m-none-eabi
+      - name: Build with no default features and alloc
+        # Use no-std target to ensure we don't link to std.
+        run: cargo build --no-default-features --features alloc --target thumbv7m-none-eabi
       - name: Test wasm
         run: wasm-pack test --headless --chrome
         if: startsWith(matrix.os, 'ubuntu')
@@ -75,7 +78,7 @@ jobs:
       matrix:
         # When updating this, the reminder to update the minimum supported
         # Rust version in Cargo.toml.
-        rust: ['1.34']
+        rust: ['1.36']
     steps:
       - uses: actions/checkout@v3
       - name: Install Rust

Cargo.toml

@@ -6,7 +6,7 @@ name = "fastrand"
 version = "1.9.0"
 authors = ["Stjepan Glavina <stjepang@gmail.com>"]
 edition = "2018"
-rust-version = "1.34"
+rust-version = "1.36"
 description = "A simple and fast random number generator"
 license = "Apache-2.0 OR MIT"
 repository = "https://github.com/smol-rs/fastrand"
@@ -16,7 +16,8 @@ exclude = ["/.*"]
 
 [features]
 default = ["std"]
-std = ["instant"]
+alloc = []
+std = ["alloc", "instant"]
 
 [target.'cfg(all(target_arch = "wasm32", not(target_os = "wasi")))'.dependencies]
 instant = { version = "0.1", optional = true }

README.md

@@ -40,6 +40,13 @@ let i = fastrand::usize(..v.len());
 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(0..20, 12);
+```
+
 Shuffle an array:
 
 ```rust

src/global_rng.rs

@@ -185,3 +185,8 @@ pub fn f32() -> f32 {
 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> {
+    with_rng(|rng| rng.choose_multiple(source, amount))
+}

src/lib.rs

@@ -29,6 +29,14 @@
 //! let elem = v[i];
 //! ```
 //!
+//! 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(0..20, 12);
+//! ```
+//!
+//!
 //! Shuffle an array:
 //!
 //! ```
@@ -76,9 +84,15 @@
 #![forbid(unsafe_code)]
 #![warn(missing_docs, missing_debug_implementations, rust_2018_idioms)]
 
+#[cfg(feature = "alloc")]
+extern crate alloc;
+
 use core::convert::{TryFrom, TryInto};
 use core::ops::{Bound, RangeBounds};
 
+#[cfg(feature = "alloc")]
+use alloc::vec::Vec;
+
 #[cfg(feature = "std")]
 #[cfg_attr(docsrs, doc(cfg(feature = "std")))]
 mod global_rng;
@@ -332,6 +346,44 @@ 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.
+    ///
+    /// The length of the returned vector equals `amount` unless the iterator
+    /// 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.
+    #[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> {
+        // Adapted from: https://docs.rs/rand/latest/rand/seq/trait.IteratorRandom.html#method.choose_multiple
+        let mut reservoir = Vec::with_capacity(amount);
+
+        reservoir.extend(source.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() {
+                let end = i + 1 + amount;
+                let k = self.usize(0..end);
+                if let Some(slot) = reservoir.get_mut(k) {
+                    *slot = elem;
+                }
+            }
+        } else {
+            // If less than one third of the `Vec` was used, reallocate
+            // so that the unused space is not wasted. There is a corner
+            // case where `amount` was much less than `self.len()`.
+            if reservoir.capacity() > 3 * reservoir.len() {
+                reservoir.shrink_to_fit();
+            }
+        }
+        reservoir
+    }
+
     rng_integer!(
         i8,
         u8,

tests/smoke.rs

@@ -119,6 +119,19 @@ fn with_seed() {
     assert_eq!(a.u64(..), b.u64(..));
 }
 
+#[test]
+fn choose_multiple() {
+    let mut a = fastrand::Rng::new();
+    let mut elements = (0..20).collect::<Vec<_>>();
+
+    while !elements.is_empty() {
+        let chosen = a.choose_multiple(0..20, 5);
+        for &x in &chosen {
+            elements.retain(|&y| y != x);
+        }
+    }
+}
+
 #[test]
 fn choice() {
     let items = [1, 4, 9, 5, 2, 3, 6, 7, 8, 0];