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
OperationFingerPrint(Operation * topOp)27
28 #[cfg(feature = "alloc")] pub mod index;
29
30 #[cfg(feature = "alloc")] use core::ops::Index;
__anonadefd66f0202(Operation *op) 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):
operator ==(const OperationFingerPrint & other) const56 /// ```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 /// ```
79 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)`.
86 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")]
114 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")]
141 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")]
168 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 /// ```
197 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`.
217 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() {
printBeforeIfEnabled__anonadefd66f0611::BasicIRPrinterConfig232 /// 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.
255 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`].
317 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")]
356 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
385 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
394 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")]
405 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")]
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")]
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
453 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
461 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
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
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 {
515 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]
525 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]
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
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
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
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
622 fn next(&mut self) -> Option<Self::Item> {
623 self.iter.next()
624 }
625
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
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
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]
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")]
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
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