1 // Copyright 2018 Developers of the Rand project.
2 //
3 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4 // https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5 // <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
6 // option. This file may not be copied, modified, or distributed
7 // except according to those terms.
8 
9 //! Sequence-related functionality
10 //!
11 //! This module provides:
12 //!
13 //! *   [`SliceRandom`] slice sampling and mutation
14 //! *   [`IteratorRandom`] iterator sampling
15 //! *   [`index::sample`] low-level API to choose multiple indices from
16 //!     `0..length`
17 //!
18 //! Also see:
19 //!
20 //! *   [`crate::distributions::weighted`] module which provides
21 //!     implementations of weighted index sampling.
22 //!
23 //! In order to make results reproducible across 32-64 bit architectures, all
24 //! `usize` indices are sampled as a `u32` where possible (also providing a
25 //! small performance boost in some cases).
26 
27 
28 #[cfg(feature = "alloc")] pub mod index;
29 
30 #[cfg(feature = "alloc")] use core::ops::Index;
31 
32 #[cfg(all(feature = "alloc", not(feature = "std")))] use crate::alloc::vec::Vec;
33 
34 #[cfg(feature = "alloc")]
35 use crate::distributions::uniform::{SampleBorrow, SampleUniform};
36 #[cfg(feature = "alloc")] use crate::distributions::WeightedError;
37 use crate::Rng;
38 
39 /// Extension trait on slices, providing random mutation and sampling methods.
40 ///
41 /// This trait is implemented on all `[T]` slice types, providing several
42 /// methods for choosing and shuffling elements. You must `use` this trait:
43 ///
44 /// ```
45 /// use rand::seq::SliceRandom;
46 ///
47 /// fn main() {
48 ///     let mut rng = rand::thread_rng();
49 ///     let mut bytes = "Hello, random!".to_string().into_bytes();
50 ///     bytes.shuffle(&mut rng);
51 ///     let str = String::from_utf8(bytes).unwrap();
52 ///     println!("{}", str);
53 /// }
54 /// ```
55 /// Example output (non-deterministic):
56 /// ```none
57 /// l,nmroHado !le
58 /// ```
59 pub trait SliceRandom {
60     /// The element type.
61     type Item;
62 
63     /// Returns a reference to one random element of the slice, or `None` if the
64     /// slice is empty.
65     ///
66     /// For slices, complexity is `O(1)`.
67     ///
68     /// # Example
69     ///
70     /// ```
71     /// use rand::thread_rng;
72     /// use rand::seq::SliceRandom;
73     ///
74     /// let choices = [1, 2, 4, 8, 16, 32];
75     /// let mut rng = thread_rng();
76     /// println!("{:?}", choices.choose(&mut rng));
77     /// assert_eq!(choices[..0].choose(&mut rng), None);
78     /// ```
choose<R>(&self, rng: &mut R) -> Option<&Self::Item> where R: Rng + ?Sized79     fn choose<R>(&self, rng: &mut R) -> Option<&Self::Item>
80     where R: Rng + ?Sized;
81 
82     /// Returns a mutable reference to one random element of the slice, or
83     /// `None` if the slice is empty.
84     ///
85     /// For slices, complexity is `O(1)`.
choose_mut<R>(&mut self, rng: &mut R) -> Option<&mut Self::Item> where R: Rng + ?Sized86     fn choose_mut<R>(&mut self, rng: &mut R) -> Option<&mut Self::Item>
87     where R: Rng + ?Sized;
88 
89     /// Chooses `amount` elements from the slice at random, without repetition,
90     /// and in random order. The returned iterator is appropriate both for
91     /// collection into a `Vec` and filling an existing buffer (see example).
92     ///
93     /// In case this API is not sufficiently flexible, use [`index::sample`].
94     ///
95     /// For slices, complexity is the same as [`index::sample`].
96     ///
97     /// # Example
98     /// ```
99     /// use rand::seq::SliceRandom;
100     ///
101     /// let mut rng = &mut rand::thread_rng();
102     /// let sample = "Hello, audience!".as_bytes();
103     ///
104     /// // collect the results into a vector:
105     /// let v: Vec<u8> = sample.choose_multiple(&mut rng, 3).cloned().collect();
106     ///
107     /// // store in a buffer:
108     /// let mut buf = [0u8; 5];
109     /// for (b, slot) in sample.choose_multiple(&mut rng, buf.len()).zip(buf.iter_mut()) {
110     ///     *slot = *b;
111     /// }
112     /// ```
113     #[cfg(feature = "alloc")]
choose_multiple<R>(&self, rng: &mut R, amount: usize) -> SliceChooseIter<Self, Self::Item> where R: Rng + ?Sized114     fn choose_multiple<R>(&self, rng: &mut R, amount: usize) -> SliceChooseIter<Self, Self::Item>
115     where R: Rng + ?Sized;
116 
117     /// Similar to [`choose`], but where the likelihood of each outcome may be
118     /// specified.
119     ///
120     /// The specified function `weight` maps each item `x` to a relative
121     /// likelihood `weight(x)`. The probability of each item being selected is
122     /// therefore `weight(x) / s`, where `s` is the sum of all `weight(x)`.
123     ///
124     /// For slices of length `n`, complexity is `O(n)`.
125     /// See also [`choose_weighted_mut`], [`distributions::weighted`].
126     ///
127     /// # Example
128     ///
129     /// ```
130     /// use rand::prelude::*;
131     ///
132     /// let choices = [('a', 2), ('b', 1), ('c', 1)];
133     /// let mut rng = thread_rng();
134     /// // 50% chance to print 'a', 25% chance to print 'b', 25% chance to print 'c'
135     /// println!("{:?}", choices.choose_weighted(&mut rng, |item| item.1).unwrap().0);
136     /// ```
137     /// [`choose`]: SliceRandom::choose
138     /// [`choose_weighted_mut`]: SliceRandom::choose_weighted_mut
139     /// [`distributions::weighted`]: crate::distributions::weighted
140     #[cfg(feature = "alloc")]
choose_weighted<R, F, B, X>( &self, rng: &mut R, weight: F, ) -> Result<&Self::Item, WeightedError> where R: Rng + ?Sized, F: Fn(&Self::Item) -> B, B: SampleBorrow<X>, X: SampleUniform + for<'a> ::core::ops::AddAssign<&'a X> + ::core::cmp::PartialOrd<X> + Clone + Default141     fn choose_weighted<R, F, B, X>(
142         &self, rng: &mut R, weight: F,
143     ) -> Result<&Self::Item, WeightedError>
144     where
145         R: Rng + ?Sized,
146         F: Fn(&Self::Item) -> B,
147         B: SampleBorrow<X>,
148         X: SampleUniform
149             + for<'a> ::core::ops::AddAssign<&'a X>
150             + ::core::cmp::PartialOrd<X>
151             + Clone
152             + Default;
153 
154     /// Similar to [`choose_mut`], but where the likelihood of each outcome may
155     /// be specified.
156     ///
157     /// The specified function `weight` maps each item `x` to a relative
158     /// likelihood `weight(x)`. The probability of each item being selected is
159     /// therefore `weight(x) / s`, where `s` is the sum of all `weight(x)`.
160     ///
161     /// For slices of length `n`, complexity is `O(n)`.
162     /// See also [`choose_weighted`], [`distributions::weighted`].
163     ///
164     /// [`choose_mut`]: SliceRandom::choose_mut
165     /// [`choose_weighted`]: SliceRandom::choose_weighted
166     /// [`distributions::weighted`]: crate::distributions::weighted
167     #[cfg(feature = "alloc")]
choose_weighted_mut<R, F, B, X>( &mut self, rng: &mut R, weight: F, ) -> Result<&mut Self::Item, WeightedError> where R: Rng + ?Sized, F: Fn(&Self::Item) -> B, B: SampleBorrow<X>, X: SampleUniform + for<'a> ::core::ops::AddAssign<&'a X> + ::core::cmp::PartialOrd<X> + Clone + Default168     fn choose_weighted_mut<R, F, B, X>(
169         &mut self, rng: &mut R, weight: F,
170     ) -> Result<&mut Self::Item, WeightedError>
171     where
172         R: Rng + ?Sized,
173         F: Fn(&Self::Item) -> B,
174         B: SampleBorrow<X>,
175         X: SampleUniform
176             + for<'a> ::core::ops::AddAssign<&'a X>
177             + ::core::cmp::PartialOrd<X>
178             + Clone
179             + Default;
180 
181     /// Shuffle a mutable slice in place.
182     ///
183     /// For slices of length `n`, complexity is `O(n)`.
184     ///
185     /// # Example
186     ///
187     /// ```
188     /// use rand::seq::SliceRandom;
189     /// use rand::thread_rng;
190     ///
191     /// let mut rng = thread_rng();
192     /// let mut y = [1, 2, 3, 4, 5];
193     /// println!("Unshuffled: {:?}", y);
194     /// y.shuffle(&mut rng);
195     /// println!("Shuffled:   {:?}", y);
196     /// ```
shuffle<R>(&mut self, rng: &mut R) where R: Rng + ?Sized197     fn shuffle<R>(&mut self, rng: &mut R)
198     where R: Rng + ?Sized;
199 
200     /// Shuffle a slice in place, but exit early.
201     ///
202     /// Returns two mutable slices from the source slice. The first contains
203     /// `amount` elements randomly permuted. The second has the remaining
204     /// elements that are not fully shuffled.
205     ///
206     /// This is an efficient method to select `amount` elements at random from
207     /// the slice, provided the slice may be mutated.
208     ///
209     /// If you only need to choose elements randomly and `amount > self.len()/2`
210     /// then you may improve performance by taking
211     /// `amount = values.len() - amount` and using only the second slice.
212     ///
213     /// If `amount` is greater than the number of elements in the slice, this
214     /// will perform a full shuffle.
215     ///
216     /// For slices, complexity is `O(m)` where `m = amount`.
partial_shuffle<R>( &mut self, rng: &mut R, amount: usize, ) -> (&mut [Self::Item], &mut [Self::Item]) where R: Rng + ?Sized217     fn partial_shuffle<R>(
218         &mut self, rng: &mut R, amount: usize,
219     ) -> (&mut [Self::Item], &mut [Self::Item])
220     where R: Rng + ?Sized;
221 }
222 
223 /// Extension trait on iterators, providing random sampling methods.
224 ///
225 /// This trait is implemented on all sized iterators, providing methods for
226 /// choosing one or more elements. You must `use` this trait:
227 ///
228 /// ```
229 /// use rand::seq::IteratorRandom;
230 ///
231 /// fn main() {
232 ///     let mut rng = rand::thread_rng();
233 ///
234 ///     let faces = "������������";
235 ///     println!("I am {}!", faces.chars().choose(&mut rng).unwrap());
236 /// }
237 /// ```
238 /// Example output (non-deterministic):
239 /// ```none
240 /// I am ��!
241 /// ```
242 pub trait IteratorRandom: Iterator + Sized {
243     /// Choose one element at random from the iterator.
244     ///
245     /// Returns `None` if and only if the iterator is empty.
246     ///
247     /// This method uses [`Iterator::size_hint`] for optimisation. With an
248     /// accurate hint and where [`Iterator::nth`] is a constant-time operation
249     /// this method can offer `O(1)` performance. Where no size hint is
250     /// available, complexity is `O(n)` where `n` is the iterator length.
251     /// Partial hints (where `lower > 0`) also improve performance.
252     ///
253     /// For slices, prefer [`SliceRandom::choose`] which guarantees `O(1)`
254     /// performance.
choose<R>(mut self, rng: &mut R) -> Option<Self::Item> where R: Rng + ?Sized255     fn choose<R>(mut self, rng: &mut R) -> Option<Self::Item>
256     where R: Rng + ?Sized {
257         let (mut lower, mut upper) = self.size_hint();
258         let mut consumed = 0;
259         let mut result = None;
260 
261         if upper == Some(lower) {
262             return if lower == 0 {
263                 None
264             } else {
265                 self.nth(gen_index(rng, lower))
266             };
267         }
268 
269         // Continue until the iterator is exhausted
270         loop {
271             if lower > 1 {
272                 let ix = gen_index(rng, lower + consumed);
273                 let skip = if ix < lower {
274                     result = self.nth(ix);
275                     lower - (ix + 1)
276                 } else {
277                     lower
278                 };
279                 if upper == Some(lower) {
280                     return result;
281                 }
282                 consumed += lower;
283                 if skip > 0 {
284                     self.nth(skip - 1);
285                 }
286             } else {
287                 let elem = self.next();
288                 if elem.is_none() {
289                     return result;
290                 }
291                 consumed += 1;
292                 let denom = consumed as f64; // accurate to 2^53 elements
293                 if rng.gen_bool(1.0 / denom) {
294                     result = elem;
295                 }
296             }
297 
298             let hint = self.size_hint();
299             lower = hint.0;
300             upper = hint.1;
301         }
302     }
303 
304     /// Collects values at random from the iterator into a supplied buffer
305     /// until that buffer is filled.
306     ///
307     /// Although the elements are selected randomly, the order of elements in
308     /// the buffer is neither stable nor fully random. If random ordering is
309     /// desired, shuffle the result.
310     ///
311     /// Returns the number of elements added to the buffer. This equals the length
312     /// of the buffer unless the iterator contains insufficient elements, in which
313     /// case this equals the number of elements available.
314     ///
315     /// Complexity is `O(n)` where `n` is the length of the iterator.
316     /// For slices, prefer [`SliceRandom::choose_multiple`].
choose_multiple_fill<R>(mut self, rng: &mut R, buf: &mut [Self::Item]) -> usize where R: Rng + ?Sized317     fn choose_multiple_fill<R>(mut self, rng: &mut R, buf: &mut [Self::Item]) -> usize
318     where R: Rng + ?Sized {
319         let amount = buf.len();
320         let mut len = 0;
321         while len < amount {
322             if let Some(elem) = self.next() {
323                 buf[len] = elem;
324                 len += 1;
325             } else {
326                 // Iterator exhausted; stop early
327                 return len;
328             }
329         }
330 
331         // Continue, since the iterator was not exhausted
332         for (i, elem) in self.enumerate() {
333             let k = gen_index(rng, i + 1 + amount);
334             if let Some(slot) = buf.get_mut(k) {
335                 *slot = elem;
336             }
337         }
338         len
339     }
340 
341     /// Collects `amount` values at random from the iterator into a vector.
342     ///
343     /// This is equivalent to `choose_multiple_fill` except for the result type.
344     ///
345     /// Although the elements are selected randomly, the order of elements in
346     /// the buffer is neither stable nor fully random. If random ordering is
347     /// desired, shuffle the result.
348     ///
349     /// The length of the returned vector equals `amount` unless the iterator
350     /// contains insufficient elements, in which case it equals the number of
351     /// elements available.
352     ///
353     /// Complexity is `O(n)` where `n` is the length of the iterator.
354     /// For slices, prefer [`SliceRandom::choose_multiple`].
355     #[cfg(feature = "alloc")]
choose_multiple<R>(mut self, rng: &mut R, amount: usize) -> Vec<Self::Item> where R: Rng + ?Sized356     fn choose_multiple<R>(mut self, rng: &mut R, amount: usize) -> Vec<Self::Item>
357     where R: Rng + ?Sized {
358         let mut reservoir = Vec::with_capacity(amount);
359         reservoir.extend(self.by_ref().take(amount));
360 
361         // Continue unless the iterator was exhausted
362         //
363         // note: this prevents iterators that "restart" from causing problems.
364         // If the iterator stops once, then so do we.
365         if reservoir.len() == amount {
366             for (i, elem) in self.enumerate() {
367                 let k = gen_index(rng, i + 1 + amount);
368                 if let Some(slot) = reservoir.get_mut(k) {
369                     *slot = elem;
370                 }
371             }
372         } else {
373             // Don't hang onto extra memory. There is a corner case where
374             // `amount` was much less than `self.len()`.
375             reservoir.shrink_to_fit();
376         }
377         reservoir
378     }
379 }
380 
381 
382 impl<T> SliceRandom for [T] {
383     type Item = T;
384 
choose<R>(&self, rng: &mut R) -> Option<&Self::Item> where R: Rng + ?Sized385     fn choose<R>(&self, rng: &mut R) -> Option<&Self::Item>
386     where R: Rng + ?Sized {
387         if self.is_empty() {
388             None
389         } else {
390             Some(&self[gen_index(rng, self.len())])
391         }
392     }
393 
choose_mut<R>(&mut self, rng: &mut R) -> Option<&mut Self::Item> where R: Rng + ?Sized394     fn choose_mut<R>(&mut self, rng: &mut R) -> Option<&mut Self::Item>
395     where R: Rng + ?Sized {
396         if self.is_empty() {
397             None
398         } else {
399             let len = self.len();
400             Some(&mut self[gen_index(rng, len)])
401         }
402     }
403 
404     #[cfg(feature = "alloc")]
choose_multiple<R>(&self, rng: &mut R, amount: usize) -> SliceChooseIter<Self, Self::Item> where R: Rng + ?Sized405     fn choose_multiple<R>(&self, rng: &mut R, amount: usize) -> SliceChooseIter<Self, Self::Item>
406     where R: Rng + ?Sized {
407         let amount = ::core::cmp::min(amount, self.len());
408         SliceChooseIter {
409             slice: self,
410             _phantom: Default::default(),
411             indices: index::sample(rng, self.len(), amount).into_iter(),
412         }
413     }
414 
415     #[cfg(feature = "alloc")]
choose_weighted<R, F, B, X>( &self, rng: &mut R, weight: F, ) -> Result<&Self::Item, WeightedError> where R: Rng + ?Sized, F: Fn(&Self::Item) -> B, B: SampleBorrow<X>, X: SampleUniform + for<'a> ::core::ops::AddAssign<&'a X> + ::core::cmp::PartialOrd<X> + Clone + Default,416     fn choose_weighted<R, F, B, X>(
417         &self, rng: &mut R, weight: F,
418     ) -> Result<&Self::Item, WeightedError>
419     where
420         R: Rng + ?Sized,
421         F: Fn(&Self::Item) -> B,
422         B: SampleBorrow<X>,
423         X: SampleUniform
424             + for<'a> ::core::ops::AddAssign<&'a X>
425             + ::core::cmp::PartialOrd<X>
426             + Clone
427             + Default,
428     {
429         use crate::distributions::{Distribution, WeightedIndex};
430         let distr = WeightedIndex::new(self.iter().map(weight))?;
431         Ok(&self[distr.sample(rng)])
432     }
433 
434     #[cfg(feature = "alloc")]
choose_weighted_mut<R, F, B, X>( &mut self, rng: &mut R, weight: F, ) -> Result<&mut Self::Item, WeightedError> where R: Rng + ?Sized, F: Fn(&Self::Item) -> B, B: SampleBorrow<X>, X: SampleUniform + for<'a> ::core::ops::AddAssign<&'a X> + ::core::cmp::PartialOrd<X> + Clone + Default,435     fn choose_weighted_mut<R, F, B, X>(
436         &mut self, rng: &mut R, weight: F,
437     ) -> Result<&mut Self::Item, WeightedError>
438     where
439         R: Rng + ?Sized,
440         F: Fn(&Self::Item) -> B,
441         B: SampleBorrow<X>,
442         X: SampleUniform
443             + for<'a> ::core::ops::AddAssign<&'a X>
444             + ::core::cmp::PartialOrd<X>
445             + Clone
446             + Default,
447     {
448         use crate::distributions::{Distribution, WeightedIndex};
449         let distr = WeightedIndex::new(self.iter().map(weight))?;
450         Ok(&mut self[distr.sample(rng)])
451     }
452 
shuffle<R>(&mut self, rng: &mut R) where R: Rng + ?Sized453     fn shuffle<R>(&mut self, rng: &mut R)
454     where R: Rng + ?Sized {
455         for i in (1..self.len()).rev() {
456             // invariant: elements with index > i have been locked in place.
457             self.swap(i, gen_index(rng, i + 1));
458         }
459     }
460 
partial_shuffle<R>( &mut self, rng: &mut R, amount: usize, ) -> (&mut [Self::Item], &mut [Self::Item]) where R: Rng + ?Sized461     fn partial_shuffle<R>(
462         &mut self, rng: &mut R, amount: usize,
463     ) -> (&mut [Self::Item], &mut [Self::Item])
464     where R: Rng + ?Sized {
465         // This applies Durstenfeld's algorithm for the
466         // [Fisher–Yates shuffle](https://en.wikipedia.org/wiki/Fisher%E2%80%93Yates_shuffle#The_modern_algorithm)
467         // for an unbiased permutation, but exits early after choosing `amount`
468         // elements.
469 
470         let len = self.len();
471         let end = if amount >= len { 0 } else { len - amount };
472 
473         for i in (end..len).rev() {
474             // invariant: elements with index > i have been locked in place.
475             self.swap(i, gen_index(rng, i + 1));
476         }
477         let r = self.split_at_mut(end);
478         (r.1, r.0)
479     }
480 }
481 
482 impl<I> IteratorRandom for I where I: Iterator + Sized {}
483 
484 
485 /// An iterator over multiple slice elements.
486 ///
487 /// This struct is created by
488 /// [`SliceRandom::choose_multiple`](trait.SliceRandom.html#tymethod.choose_multiple).
489 #[cfg(feature = "alloc")]
490 #[derive(Debug)]
491 pub struct SliceChooseIter<'a, S: ?Sized + 'a, T: 'a> {
492     slice: &'a S,
493     _phantom: ::core::marker::PhantomData<T>,
494     indices: index::IndexVecIntoIter,
495 }
496 
497 #[cfg(feature = "alloc")]
498 impl<'a, S: Index<usize, Output = T> + ?Sized + 'a, T: 'a> Iterator for SliceChooseIter<'a, S, T> {
499     type Item = &'a T;
500 
next(&mut self) -> Option<Self::Item>501     fn next(&mut self) -> Option<Self::Item> {
502         // TODO: investigate using SliceIndex::get_unchecked when stable
503         self.indices.next().map(|i| &self.slice[i as usize])
504     }
505 
size_hint(&self) -> (usize, Option<usize>)506     fn size_hint(&self) -> (usize, Option<usize>) {
507         (self.indices.len(), Some(self.indices.len()))
508     }
509 }
510 
511 #[cfg(feature = "alloc")]
512 impl<'a, S: Index<usize, Output = T> + ?Sized + 'a, T: 'a> ExactSizeIterator
513     for SliceChooseIter<'a, S, T>
514 {
len(&self) -> usize515     fn len(&self) -> usize {
516         self.indices.len()
517     }
518 }
519 
520 
521 // Sample a number uniformly between 0 and `ubound`. Uses 32-bit sampling where
522 // possible, primarily in order to produce the same output on 32-bit and 64-bit
523 // platforms.
524 #[inline]
gen_index<R: Rng + ?Sized>(rng: &mut R, ubound: usize) -> usize525 fn gen_index<R: Rng + ?Sized>(rng: &mut R, ubound: usize) -> usize {
526     if ubound <= (core::u32::MAX as usize) {
527         rng.gen_range(0, ubound as u32) as usize
528     } else {
529         rng.gen_range(0, ubound)
530     }
531 }
532 
533 
534 #[cfg(test)]
535 mod test {
536     use super::*;
537     #[cfg(feature = "alloc")] use crate::Rng;
538     #[cfg(all(feature = "alloc", not(feature = "std")))] use alloc::vec::Vec;
539 
540     #[test]
test_slice_choose()541     fn test_slice_choose() {
542         let mut r = crate::test::rng(107);
543         let chars = [
544             'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n',
545         ];
546         let mut chosen = [0i32; 14];
547         // The below all use a binomial distribution with n=1000, p=1/14.
548         // binocdf(40, 1000, 1/14) ~= 2e-5; 1-binocdf(106, ..) ~= 2e-5
549         for _ in 0..1000 {
550             let picked = *chars.choose(&mut r).unwrap();
551             chosen[(picked as usize) - ('a' as usize)] += 1;
552         }
553         for count in chosen.iter() {
554             assert!(40 < *count && *count < 106);
555         }
556 
557         chosen.iter_mut().for_each(|x| *x = 0);
558         for _ in 0..1000 {
559             *chosen.choose_mut(&mut r).unwrap() += 1;
560         }
561         for count in chosen.iter() {
562             assert!(40 < *count && *count < 106);
563         }
564 
565         let mut v: [isize; 0] = [];
566         assert_eq!(v.choose(&mut r), None);
567         assert_eq!(v.choose_mut(&mut r), None);
568     }
569 
570     #[derive(Clone)]
571     struct UnhintedIterator<I: Iterator + Clone> {
572         iter: I,
573     }
574     impl<I: Iterator + Clone> Iterator for UnhintedIterator<I> {
575         type Item = I::Item;
576 
next(&mut self) -> Option<Self::Item>577         fn next(&mut self) -> Option<Self::Item> {
578             self.iter.next()
579         }
580     }
581 
582     #[derive(Clone)]
583     struct ChunkHintedIterator<I: ExactSizeIterator + Iterator + Clone> {
584         iter: I,
585         chunk_remaining: usize,
586         chunk_size: usize,
587         hint_total_size: bool,
588     }
589     impl<I: ExactSizeIterator + Iterator + Clone> Iterator for ChunkHintedIterator<I> {
590         type Item = I::Item;
591 
next(&mut self) -> Option<Self::Item>592         fn next(&mut self) -> Option<Self::Item> {
593             if self.chunk_remaining == 0 {
594                 self.chunk_remaining = ::core::cmp::min(self.chunk_size, self.iter.len());
595             }
596             self.chunk_remaining = self.chunk_remaining.saturating_sub(1);
597 
598             self.iter.next()
599         }
600 
size_hint(&self) -> (usize, Option<usize>)601         fn size_hint(&self) -> (usize, Option<usize>) {
602             (
603                 self.chunk_remaining,
604                 if self.hint_total_size {
605                     Some(self.iter.len())
606                 } else {
607                     None
608                 },
609             )
610         }
611     }
612 
613     #[derive(Clone)]
614     struct WindowHintedIterator<I: ExactSizeIterator + Iterator + Clone> {
615         iter: I,
616         window_size: usize,
617         hint_total_size: bool,
618     }
619     impl<I: ExactSizeIterator + Iterator + Clone> Iterator for WindowHintedIterator<I> {
620         type Item = I::Item;
621 
next(&mut self) -> Option<Self::Item>622         fn next(&mut self) -> Option<Self::Item> {
623             self.iter.next()
624         }
625 
size_hint(&self) -> (usize, Option<usize>)626         fn size_hint(&self) -> (usize, Option<usize>) {
627             (
628                 ::core::cmp::min(self.iter.len(), self.window_size),
629                 if self.hint_total_size {
630                     Some(self.iter.len())
631                 } else {
632                     None
633                 },
634             )
635         }
636     }
637 
638     #[test]
639     #[cfg_attr(miri, ignore)] // Miri is too slow
test_iterator_choose()640     fn test_iterator_choose() {
641         let r = &mut crate::test::rng(109);
642         fn test_iter<R: Rng + ?Sized, Iter: Iterator<Item = usize> + Clone>(r: &mut R, iter: Iter) {
643             let mut chosen = [0i32; 9];
644             for _ in 0..1000 {
645                 let picked = iter.clone().choose(r).unwrap();
646                 chosen[picked] += 1;
647             }
648             for count in chosen.iter() {
649                 // Samples should follow Binomial(1000, 1/9)
650                 // Octave: binopdf(x, 1000, 1/9) gives the prob of *count == x
651                 // Note: have seen 153, which is unlikely but not impossible.
652                 assert!(
653                     72 < *count && *count < 154,
654                     "count not close to 1000/9: {}",
655                     count
656                 );
657             }
658         }
659 
660         test_iter(r, 0..9);
661         test_iter(r, [0, 1, 2, 3, 4, 5, 6, 7, 8].iter().cloned());
662         #[cfg(feature = "alloc")]
663         test_iter(r, (0..9).collect::<Vec<_>>().into_iter());
664         test_iter(r, UnhintedIterator { iter: 0..9 });
665         test_iter(r, ChunkHintedIterator {
666             iter: 0..9,
667             chunk_size: 4,
668             chunk_remaining: 4,
669             hint_total_size: false,
670         });
671         test_iter(r, ChunkHintedIterator {
672             iter: 0..9,
673             chunk_size: 4,
674             chunk_remaining: 4,
675             hint_total_size: true,
676         });
677         test_iter(r, WindowHintedIterator {
678             iter: 0..9,
679             window_size: 2,
680             hint_total_size: false,
681         });
682         test_iter(r, WindowHintedIterator {
683             iter: 0..9,
684             window_size: 2,
685             hint_total_size: true,
686         });
687 
688         assert_eq!((0..0).choose(r), None);
689         assert_eq!(UnhintedIterator { iter: 0..0 }.choose(r), None);
690     }
691 
692     #[test]
693     #[cfg_attr(miri, ignore)] // Miri is too slow
test_shuffle()694     fn test_shuffle() {
695         let mut r = crate::test::rng(108);
696         let empty: &mut [isize] = &mut [];
697         empty.shuffle(&mut r);
698         let mut one = [1];
699         one.shuffle(&mut r);
700         let b: &[_] = &[1];
701         assert_eq!(one, b);
702 
703         let mut two = [1, 2];
704         two.shuffle(&mut r);
705         assert!(two == [1, 2] || two == [2, 1]);
706 
707         fn move_last(slice: &mut [usize], pos: usize) {
708             // use slice[pos..].rotate_left(1); once we can use that
709             let last_val = slice[pos];
710             for i in pos..slice.len() - 1 {
711                 slice[i] = slice[i + 1];
712             }
713             *slice.last_mut().unwrap() = last_val;
714         }
715         let mut counts = [0i32; 24];
716         for _ in 0..10000 {
717             let mut arr: [usize; 4] = [0, 1, 2, 3];
718             arr.shuffle(&mut r);
719             let mut permutation = 0usize;
720             let mut pos_value = counts.len();
721             for i in 0..4 {
722                 pos_value /= 4 - i;
723                 let pos = arr.iter().position(|&x| x == i).unwrap();
724                 assert!(pos < (4 - i));
725                 permutation += pos * pos_value;
726                 move_last(&mut arr, pos);
727                 assert_eq!(arr[3], i);
728             }
729             for i in 0..4 {
730                 assert_eq!(arr[i], i);
731             }
732             counts[permutation] += 1;
733         }
734         for count in counts.iter() {
735             // Binomial(10000, 1/24) with average 416.667
736             // Octave: binocdf(n, 10000, 1/24)
737             // 99.9% chance samples lie within this range:
738             assert!(352 <= *count && *count <= 483, "count: {}", count);
739         }
740     }
741 
742     #[test]
test_partial_shuffle()743     fn test_partial_shuffle() {
744         let mut r = crate::test::rng(118);
745 
746         let mut empty: [u32; 0] = [];
747         let res = empty.partial_shuffle(&mut r, 10);
748         assert_eq!((res.0.len(), res.1.len()), (0, 0));
749 
750         let mut v = [1, 2, 3, 4, 5];
751         let res = v.partial_shuffle(&mut r, 2);
752         assert_eq!((res.0.len(), res.1.len()), (2, 3));
753         assert!(res.0[0] != res.0[1]);
754         // First elements are only modified if selected, so at least one isn't modified:
755         assert!(res.1[0] == 1 || res.1[1] == 2 || res.1[2] == 3);
756     }
757 
758     #[test]
759     #[cfg(feature = "alloc")]
test_sample_iter()760     fn test_sample_iter() {
761         let min_val = 1;
762         let max_val = 100;
763 
764         let mut r = crate::test::rng(401);
765         let vals = (min_val..max_val).collect::<Vec<i32>>();
766         let small_sample = vals.iter().choose_multiple(&mut r, 5);
767         let large_sample = vals.iter().choose_multiple(&mut r, vals.len() + 5);
768 
769         assert_eq!(small_sample.len(), 5);
770         assert_eq!(large_sample.len(), vals.len());
771         // no randomization happens when amount >= len
772         assert_eq!(large_sample, vals.iter().collect::<Vec<_>>());
773 
774         assert!(small_sample
775             .iter()
776             .all(|e| { **e >= min_val && **e <= max_val }));
777     }
778 
779     #[test]
780     #[cfg(feature = "alloc")]
781     #[cfg_attr(miri, ignore)] // Miri is too slow
test_weighted()782     fn test_weighted() {
783         let mut r = crate::test::rng(406);
784         const N_REPS: u32 = 3000;
785         let weights = [1u32, 2, 3, 0, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7];
786         let total_weight = weights.iter().sum::<u32>() as f32;
787 
788         let verify = |result: [i32; 14]| {
789             for (i, count) in result.iter().enumerate() {
790                 let exp = (weights[i] * N_REPS) as f32 / total_weight;
791                 let mut err = (*count as f32 - exp).abs();
792                 if err != 0.0 {
793                     err /= exp;
794                 }
795                 assert!(err <= 0.25);
796             }
797         };
798 
799         // choose_weighted
800         fn get_weight<T>(item: &(u32, T)) -> u32 {
801             item.0
802         }
803         let mut chosen = [0i32; 14];
804         let mut items = [(0u32, 0usize); 14]; // (weight, index)
805         for (i, item) in items.iter_mut().enumerate() {
806             *item = (weights[i], i);
807         }
808         for _ in 0..N_REPS {
809             let item = items.choose_weighted(&mut r, get_weight).unwrap();
810             chosen[item.1] += 1;
811         }
812         verify(chosen);
813 
814         // choose_weighted_mut
815         let mut items = [(0u32, 0i32); 14]; // (weight, count)
816         for (i, item) in items.iter_mut().enumerate() {
817             *item = (weights[i], 0);
818         }
819         for _ in 0..N_REPS {
820             items.choose_weighted_mut(&mut r, get_weight).unwrap().1 += 1;
821         }
822         for (ch, item) in chosen.iter_mut().zip(items.iter()) {
823             *ch = item.1;
824         }
825         verify(chosen);
826 
827         // Check error cases
828         let empty_slice = &mut [10][0..0];
829         assert_eq!(
830             empty_slice.choose_weighted(&mut r, |_| 1),
831             Err(WeightedError::NoItem)
832         );
833         assert_eq!(
834             empty_slice.choose_weighted_mut(&mut r, |_| 1),
835             Err(WeightedError::NoItem)
836         );
837         assert_eq!(
838             ['x'].choose_weighted_mut(&mut r, |_| 0),
839             Err(WeightedError::AllWeightsZero)
840         );
841         assert_eq!(
842             [0, -1].choose_weighted_mut(&mut r, |x| *x),
843             Err(WeightedError::InvalidWeight)
844         );
845         assert_eq!(
846             [-1, 0].choose_weighted_mut(&mut r, |x| *x),
847             Err(WeightedError::InvalidWeight)
848         );
849     }
850 }
851