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