#[cfg(feature = "alloc")]
#[cfg_attr(doc_cfg, doc(cfg(feature = "alloc")))]
pub mod index;
#[cfg(feature = "alloc")] use core::ops::Index;
#[cfg(feature = "alloc")] use alloc::vec::Vec;
#[cfg(feature = "alloc")]
use crate::distributions::uniform::{SampleBorrow, SampleUniform};
#[cfg(feature = "alloc")] use crate::distributions::WeightedError;
use crate::Rng;
pub trait SliceRandom {
type Item;
fn choose<R>(&self, rng: &mut R) -> Option<&Self::Item>
where R: Rng + ?Sized;
fn choose_mut<R>(&mut self, rng: &mut R) -> Option<&mut Self::Item>
where R: Rng + ?Sized;
#[cfg(feature = "alloc")]
#[cfg_attr(doc_cfg, doc(cfg(feature = "alloc")))]
fn choose_multiple<R>(&self, rng: &mut R, amount: usize) -> SliceChooseIter<Self, Self::Item>
where R: Rng + ?Sized;
#[cfg(feature = "alloc")]
#[cfg_attr(doc_cfg, doc(cfg(feature = "alloc")))]
fn 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;
#[cfg(feature = "alloc")]
#[cfg_attr(doc_cfg, doc(cfg(feature = "alloc")))]
fn 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;
#[cfg(feature = "std")]
#[cfg_attr(doc_cfg, doc(cfg(feature = "std")))]
fn choose_multiple_weighted<R, F, X>(
&self, rng: &mut R, amount: usize, weight: F,
) -> Result<SliceChooseIter<Self, Self::Item>, WeightedError>
where
R: Rng + ?Sized,
F: Fn(&Self::Item) -> X,
X: Into<f64>;
fn shuffle<R>(&mut self, rng: &mut R)
where R: Rng + ?Sized;
fn partial_shuffle<R>(
&mut self, rng: &mut R, amount: usize,
) -> (&mut [Self::Item], &mut [Self::Item])
where R: Rng + ?Sized;
}
pub trait IteratorRandom: Iterator + Sized {
fn choose<R>(mut self, rng: &mut R) -> Option<Self::Item>
where R: Rng + ?Sized {
let (mut lower, mut upper) = self.size_hint();
let mut consumed = 0;
let mut result = None;
if upper == Some(lower) {
return if lower == 0 {
None
} else {
self.nth(gen_index(rng, lower))
};
}
loop {
if lower > 1 {
let ix = gen_index(rng, lower + consumed);
let skip = if ix < lower {
result = self.nth(ix);
lower - (ix + 1)
} else {
lower
};
if upper == Some(lower) {
return result;
}
consumed += lower;
if skip > 0 {
self.nth(skip - 1);
}
} else {
let elem = self.next();
if elem.is_none() {
return result;
}
consumed += 1;
if gen_index(rng, consumed) == 0 {
result = elem;
}
}
let hint = self.size_hint();
lower = hint.0;
upper = hint.1;
}
}
fn choose_stable<R>(mut self, rng: &mut R) -> Option<Self::Item>
where R: Rng + ?Sized {
let mut consumed = 0;
let mut result = None;
loop {
let mut next = 0;
let (lower, _) = self.size_hint();
if lower >= 2 {
let highest_selected = (0..lower)
.filter(|ix| gen_index(rng, consumed+ix+1) == 0)
.last();
consumed += lower;
next = lower;
if let Some(ix) = highest_selected {
result = self.nth(ix);
next -= ix + 1;
debug_assert!(result.is_some(), "iterator shorter than size_hint().0");
}
}
let elem = self.nth(next);
if elem.is_none() {
return result
}
if gen_index(rng, consumed+1) == 0 {
result = elem;
}
consumed += 1;
}
}
fn choose_multiple_fill<R>(mut self, rng: &mut R, buf: &mut [Self::Item]) -> usize
where R: Rng + ?Sized {
let amount = buf.len();
let mut len = 0;
while len < amount {
if let Some(elem) = self.next() {
buf[len] = elem;
len += 1;
} else {
return len;
}
}
for (i, elem) in self.enumerate() {
let k = gen_index(rng, i + 1 + amount);
if let Some(slot) = buf.get_mut(k) {
*slot = elem;
}
}
len
}
#[cfg(feature = "alloc")]
#[cfg_attr(doc_cfg, doc(cfg(feature = "alloc")))]
fn choose_multiple<R>(mut self, rng: &mut R, amount: usize) -> Vec<Self::Item>
where R: Rng + ?Sized {
let mut reservoir = Vec::with_capacity(amount);
reservoir.extend(self.by_ref().take(amount));
if reservoir.len() == amount {
for (i, elem) in self.enumerate() {
let k = gen_index(rng, i + 1 + amount);
if let Some(slot) = reservoir.get_mut(k) {
*slot = elem;
}
}
} else {
reservoir.shrink_to_fit();
}
reservoir
}
}
impl<T> SliceRandom for [T] {
type Item = T;
fn choose<R>(&self, rng: &mut R) -> Option<&Self::Item>
where R: Rng + ?Sized {
if self.is_empty() {
None
} else {
Some(&self[gen_index(rng, self.len())])
}
}
fn choose_mut<R>(&mut self, rng: &mut R) -> Option<&mut Self::Item>
where R: Rng + ?Sized {
if self.is_empty() {
None
} else {
let len = self.len();
Some(&mut self[gen_index(rng, len)])
}
}
#[cfg(feature = "alloc")]
fn choose_multiple<R>(&self, rng: &mut R, amount: usize) -> SliceChooseIter<Self, Self::Item>
where R: Rng + ?Sized {
let amount = ::core::cmp::min(amount, self.len());
SliceChooseIter {
slice: self,
_phantom: Default::default(),
indices: index::sample(rng, self.len(), amount).into_iter(),
}
}
#[cfg(feature = "alloc")]
fn 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,
{
use crate::distributions::{Distribution, WeightedIndex};
let distr = WeightedIndex::new(self.iter().map(weight))?;
Ok(&self[distr.sample(rng)])
}
#[cfg(feature = "alloc")]
fn 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,
{
use crate::distributions::{Distribution, WeightedIndex};
let distr = WeightedIndex::new(self.iter().map(weight))?;
Ok(&mut self[distr.sample(rng)])
}
#[cfg(feature = "std")]
fn choose_multiple_weighted<R, F, X>(
&self, rng: &mut R, amount: usize, weight: F,
) -> Result<SliceChooseIter<Self, Self::Item>, WeightedError>
where
R: Rng + ?Sized,
F: Fn(&Self::Item) -> X,
X: Into<f64>,
{
let amount = ::core::cmp::min(amount, self.len());
Ok(SliceChooseIter {
slice: self,
_phantom: Default::default(),
indices: index::sample_weighted(
rng,
self.len(),
|idx| weight(&self[idx]).into(),
amount,
)?
.into_iter(),
})
}
fn shuffle<R>(&mut self, rng: &mut R)
where R: Rng + ?Sized {
for i in (1..self.len()).rev() {
self.swap(i, gen_index(rng, i + 1));
}
}
fn partial_shuffle<R>(
&mut self, rng: &mut R, amount: usize,
) -> (&mut [Self::Item], &mut [Self::Item])
where R: Rng + ?Sized {
let len = self.len();
let end = if amount >= len { 0 } else { len - amount };
for i in (end..len).rev() {
self.swap(i, gen_index(rng, i + 1));
}
let r = self.split_at_mut(end);
(r.1, r.0)
}
}
impl<I> IteratorRandom for I where I: Iterator + Sized {}
#[cfg(feature = "alloc")]
#[cfg_attr(doc_cfg, doc(cfg(feature = "alloc")))]
#[derive(Debug)]
pub struct SliceChooseIter<'a, S: ?Sized + 'a, T: 'a> {
slice: &'a S,
_phantom: ::core::marker::PhantomData<T>,
indices: index::IndexVecIntoIter,
}
#[cfg(feature = "alloc")]
impl<'a, S: Index<usize, Output = T> + ?Sized + 'a, T: 'a> Iterator for SliceChooseIter<'a, S, T> {
type Item = &'a T;
fn next(&mut self) -> Option<Self::Item> {
self.indices.next().map(|i| &self.slice[i as usize])
}
fn size_hint(&self) -> (usize, Option<usize>) {
(self.indices.len(), Some(self.indices.len()))
}
}
#[cfg(feature = "alloc")]
impl<'a, S: Index<usize, Output = T> + ?Sized + 'a, T: 'a> ExactSizeIterator
for SliceChooseIter<'a, S, T>
{
fn len(&self) -> usize {
self.indices.len()
}
}
#[inline]
fn gen_index<R: Rng + ?Sized>(rng: &mut R, ubound: usize) -> usize {
if ubound <= (core::u32::MAX as usize) {
rng.gen_range(0..ubound as u32) as usize
} else {
rng.gen_range(0..ubound)
}
}
#[cfg(test)]
mod test {
use super::*;
#[cfg(feature = "alloc")] use crate::Rng;
#[cfg(all(feature = "alloc", not(feature = "std")))] use alloc::vec::Vec;
#[test]
fn test_slice_choose() {
let mut r = crate::test::rng(107);
let chars = [
'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n',
];
let mut chosen = [0i32; 14];
for _ in 0..1000 {
let picked = *chars.choose(&mut r).unwrap();
chosen[(picked as usize) - ('a' as usize)] += 1;
}
for count in chosen.iter() {
assert!(40 < *count && *count < 106);
}
chosen.iter_mut().for_each(|x| *x = 0);
for _ in 0..1000 {
*chosen.choose_mut(&mut r).unwrap() += 1;
}
for count in chosen.iter() {
assert!(40 < *count && *count < 106);
}
let mut v: [isize; 0] = [];
assert_eq!(v.choose(&mut r), None);
assert_eq!(v.choose_mut(&mut r), None);
}
#[test]
fn value_stability_slice() {
let mut r = crate::test::rng(413);
let chars = [
'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n',
];
let mut nums = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12];
assert_eq!(chars.choose(&mut r), Some(&'l'));
assert_eq!(nums.choose_mut(&mut r), Some(&mut 10));
#[cfg(feature = "alloc")]
assert_eq!(
&chars
.choose_multiple(&mut r, 8)
.cloned()
.collect::<Vec<char>>(),
&['d', 'm', 'b', 'n', 'c', 'k', 'h', 'e']
);
#[cfg(feature = "alloc")]
assert_eq!(chars.choose_weighted(&mut r, |_| 1), Ok(&'f'));
#[cfg(feature = "alloc")]
assert_eq!(nums.choose_weighted_mut(&mut r, |_| 1), Ok(&mut 5));
let mut r = crate::test::rng(414);
nums.shuffle(&mut r);
assert_eq!(nums, [9, 5, 3, 10, 7, 12, 8, 11, 6, 4, 0, 2, 1]);
nums = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12];
let res = nums.partial_shuffle(&mut r, 6);
assert_eq!(res.0, &mut [7, 4, 8, 6, 9, 3]);
assert_eq!(res.1, &mut [0, 1, 2, 12, 11, 5, 10]);
}
#[derive(Clone)]
struct UnhintedIterator<I: Iterator + Clone> {
iter: I,
}
impl<I: Iterator + Clone> Iterator for UnhintedIterator<I> {
type Item = I::Item;
fn next(&mut self) -> Option<Self::Item> {
self.iter.next()
}
}
#[derive(Clone)]
struct ChunkHintedIterator<I: ExactSizeIterator + Iterator + Clone> {
iter: I,
chunk_remaining: usize,
chunk_size: usize,
hint_total_size: bool,
}
impl<I: ExactSizeIterator + Iterator + Clone> Iterator for ChunkHintedIterator<I> {
type Item = I::Item;
fn next(&mut self) -> Option<Self::Item> {
if self.chunk_remaining == 0 {
self.chunk_remaining = ::core::cmp::min(self.chunk_size, self.iter.len());
}
self.chunk_remaining = self.chunk_remaining.saturating_sub(1);
self.iter.next()
}
fn size_hint(&self) -> (usize, Option<usize>) {
(
self.chunk_remaining,
if self.hint_total_size {
Some(self.iter.len())
} else {
None
},
)
}
}
#[derive(Clone)]
struct WindowHintedIterator<I: ExactSizeIterator + Iterator + Clone> {
iter: I,
window_size: usize,
hint_total_size: bool,
}
impl<I: ExactSizeIterator + Iterator + Clone> Iterator for WindowHintedIterator<I> {
type Item = I::Item;
fn next(&mut self) -> Option<Self::Item> {
self.iter.next()
}
fn size_hint(&self) -> (usize, Option<usize>) {
(
::core::cmp::min(self.iter.len(), self.window_size),
if self.hint_total_size {
Some(self.iter.len())
} else {
None
},
)
}
}
#[test]
#[cfg_attr(miri, ignore)]
fn test_iterator_choose() {
let r = &mut crate::test::rng(109);
fn test_iter<R: Rng + ?Sized, Iter: Iterator<Item = usize> + Clone>(r: &mut R, iter: Iter) {
let mut chosen = [0i32; 9];
for _ in 0..1000 {
let picked = iter.clone().choose(r).unwrap();
chosen[picked] += 1;
}
for count in chosen.iter() {
assert!(
72 < *count && *count < 154,
"count not close to 1000/9: {}",
count
);
}
}
test_iter(r, 0..9);
test_iter(r, [0, 1, 2, 3, 4, 5, 6, 7, 8].iter().cloned());
#[cfg(feature = "alloc")]
test_iter(r, (0..9).collect::<Vec<_>>().into_iter());
test_iter(r, UnhintedIterator { iter: 0..9 });
test_iter(r, ChunkHintedIterator {
iter: 0..9,
chunk_size: 4,
chunk_remaining: 4,
hint_total_size: false,
});
test_iter(r, ChunkHintedIterator {
iter: 0..9,
chunk_size: 4,
chunk_remaining: 4,
hint_total_size: true,
});
test_iter(r, WindowHintedIterator {
iter: 0..9,
window_size: 2,
hint_total_size: false,
});
test_iter(r, WindowHintedIterator {
iter: 0..9,
window_size: 2,
hint_total_size: true,
});
assert_eq!((0..0).choose(r), None);
assert_eq!(UnhintedIterator { iter: 0..0 }.choose(r), None);
}
#[test]
#[cfg_attr(miri, ignore)]
fn test_iterator_choose_stable() {
let r = &mut crate::test::rng(109);
fn test_iter<R: Rng + ?Sized, Iter: Iterator<Item = usize> + Clone>(r: &mut R, iter: Iter) {
let mut chosen = [0i32; 9];
for _ in 0..1000 {
let picked = iter.clone().choose_stable(r).unwrap();
chosen[picked] += 1;
}
for count in chosen.iter() {
assert!(
72 < *count && *count < 154,
"count not close to 1000/9: {}",
count
);
}
}
test_iter(r, 0..9);
test_iter(r, [0, 1, 2, 3, 4, 5, 6, 7, 8].iter().cloned());
#[cfg(feature = "alloc")]
test_iter(r, (0..9).collect::<Vec<_>>().into_iter());
test_iter(r, UnhintedIterator { iter: 0..9 });
test_iter(r, ChunkHintedIterator {
iter: 0..9,
chunk_size: 4,
chunk_remaining: 4,
hint_total_size: false,
});
test_iter(r, ChunkHintedIterator {
iter: 0..9,
chunk_size: 4,
chunk_remaining: 4,
hint_total_size: true,
});
test_iter(r, WindowHintedIterator {
iter: 0..9,
window_size: 2,
hint_total_size: false,
});
test_iter(r, WindowHintedIterator {
iter: 0..9,
window_size: 2,
hint_total_size: true,
});
assert_eq!((0..0).choose(r), None);
assert_eq!(UnhintedIterator { iter: 0..0 }.choose(r), None);
}
#[test]
#[cfg_attr(miri, ignore)]
fn test_iterator_choose_stable_stability() {
fn test_iter(iter: impl Iterator<Item = usize> + Clone) -> [i32; 9] {
let r = &mut crate::test::rng(109);
let mut chosen = [0i32; 9];
for _ in 0..1000 {
let picked = iter.clone().choose_stable(r).unwrap();
chosen[picked] += 1;
}
chosen
}
let reference = test_iter(0..9);
assert_eq!(test_iter([0, 1, 2, 3, 4, 5, 6, 7, 8].iter().cloned()), reference);
#[cfg(feature = "alloc")]
assert_eq!(test_iter((0..9).collect::<Vec<_>>().into_iter()), reference);
assert_eq!(test_iter(UnhintedIterator { iter: 0..9 }), reference);
assert_eq!(test_iter(ChunkHintedIterator {
iter: 0..9,
chunk_size: 4,
chunk_remaining: 4,
hint_total_size: false,
}), reference);
assert_eq!(test_iter(ChunkHintedIterator {
iter: 0..9,
chunk_size: 4,
chunk_remaining: 4,
hint_total_size: true,
}), reference);
assert_eq!(test_iter(WindowHintedIterator {
iter: 0..9,
window_size: 2,
hint_total_size: false,
}), reference);
assert_eq!(test_iter(WindowHintedIterator {
iter: 0..9,
window_size: 2,
hint_total_size: true,
}), reference);
}
#[test]
#[cfg_attr(miri, ignore)]
fn test_shuffle() {
let mut r = crate::test::rng(108);
let empty: &mut [isize] = &mut [];
empty.shuffle(&mut r);
let mut one = [1];
one.shuffle(&mut r);
let b: &[_] = &[1];
assert_eq!(one, b);
let mut two = [1, 2];
two.shuffle(&mut r);
assert!(two == [1, 2] || two == [2, 1]);
fn move_last(slice: &mut [usize], pos: usize) {
let last_val = slice[pos];
for i in pos..slice.len() - 1 {
slice[i] = slice[i + 1];
}
*slice.last_mut().unwrap() = last_val;
}
let mut counts = [0i32; 24];
for _ in 0..10000 {
let mut arr: [usize; 4] = [0, 1, 2, 3];
arr.shuffle(&mut r);
let mut permutation = 0usize;
let mut pos_value = counts.len();
for i in 0..4 {
pos_value /= 4 - i;
let pos = arr.iter().position(|&x| x == i).unwrap();
assert!(pos < (4 - i));
permutation += pos * pos_value;
move_last(&mut arr, pos);
assert_eq!(arr[3], i);
}
for i in 0..4 {
assert_eq!(arr[i], i);
}
counts[permutation] += 1;
}
for count in counts.iter() {
assert!(352 <= *count && *count <= 483, "count: {}", count);
}
}
#[test]
fn test_partial_shuffle() {
let mut r = crate::test::rng(118);
let mut empty: [u32; 0] = [];
let res = empty.partial_shuffle(&mut r, 10);
assert_eq!((res.0.len(), res.1.len()), (0, 0));
let mut v = [1, 2, 3, 4, 5];
let res = v.partial_shuffle(&mut r, 2);
assert_eq!((res.0.len(), res.1.len()), (2, 3));
assert!(res.0[0] != res.0[1]);
assert!(res.1[0] == 1 || res.1[1] == 2 || res.1[2] == 3);
}
#[test]
#[cfg(feature = "alloc")]
fn test_sample_iter() {
let min_val = 1;
let max_val = 100;
let mut r = crate::test::rng(401);
let vals = (min_val..max_val).collect::<Vec<i32>>();
let small_sample = vals.iter().choose_multiple(&mut r, 5);
let large_sample = vals.iter().choose_multiple(&mut r, vals.len() + 5);
assert_eq!(small_sample.len(), 5);
assert_eq!(large_sample.len(), vals.len());
assert_eq!(large_sample, vals.iter().collect::<Vec<_>>());
assert!(small_sample
.iter()
.all(|e| { **e >= min_val && **e <= max_val }));
}
#[test]
#[cfg(feature = "alloc")]
#[cfg_attr(miri, ignore)]
fn test_weighted() {
let mut r = crate::test::rng(406);
const N_REPS: u32 = 3000;
let weights = [1u32, 2, 3, 0, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7];
let total_weight = weights.iter().sum::<u32>() as f32;
let verify = |result: [i32; 14]| {
for (i, count) in result.iter().enumerate() {
let exp = (weights[i] * N_REPS) as f32 / total_weight;
let mut err = (*count as f32 - exp).abs();
if err != 0.0 {
err /= exp;
}
assert!(err <= 0.25);
}
};
fn get_weight<T>(item: &(u32, T)) -> u32 {
item.0
}
let mut chosen = [0i32; 14];
let mut items = [(0u32, 0usize); 14];
for (i, item) in items.iter_mut().enumerate() {
*item = (weights[i], i);
}
for _ in 0..N_REPS {
let item = items.choose_weighted(&mut r, get_weight).unwrap();
chosen[item.1] += 1;
}
verify(chosen);
let mut items = [(0u32, 0i32); 14];
for (i, item) in items.iter_mut().enumerate() {
*item = (weights[i], 0);
}
for _ in 0..N_REPS {
items.choose_weighted_mut(&mut r, get_weight).unwrap().1 += 1;
}
for (ch, item) in chosen.iter_mut().zip(items.iter()) {
*ch = item.1;
}
verify(chosen);
let empty_slice = &mut [10][0..0];
assert_eq!(
empty_slice.choose_weighted(&mut r, |_| 1),
Err(WeightedError::NoItem)
);
assert_eq!(
empty_slice.choose_weighted_mut(&mut r, |_| 1),
Err(WeightedError::NoItem)
);
assert_eq!(
['x'].choose_weighted_mut(&mut r, |_| 0),
Err(WeightedError::AllWeightsZero)
);
assert_eq!(
[0, -1].choose_weighted_mut(&mut r, |x| *x),
Err(WeightedError::InvalidWeight)
);
assert_eq!(
[-1, 0].choose_weighted_mut(&mut r, |x| *x),
Err(WeightedError::InvalidWeight)
);
}
#[test]
fn value_stability_choose() {
fn choose<I: Iterator<Item = u32>>(iter: I) -> Option<u32> {
let mut rng = crate::test::rng(411);
iter.choose(&mut rng)
}
assert_eq!(choose([].iter().cloned()), None);
assert_eq!(choose(0..100), Some(33));
assert_eq!(choose(UnhintedIterator { iter: 0..100 }), Some(40));
assert_eq!(
choose(ChunkHintedIterator {
iter: 0..100,
chunk_size: 32,
chunk_remaining: 32,
hint_total_size: false,
}),
Some(39)
);
assert_eq!(
choose(ChunkHintedIterator {
iter: 0..100,
chunk_size: 32,
chunk_remaining: 32,
hint_total_size: true,
}),
Some(39)
);
assert_eq!(
choose(WindowHintedIterator {
iter: 0..100,
window_size: 32,
hint_total_size: false,
}),
Some(90)
);
assert_eq!(
choose(WindowHintedIterator {
iter: 0..100,
window_size: 32,
hint_total_size: true,
}),
Some(90)
);
}
#[test]
fn value_stability_choose_stable() {
fn choose<I: Iterator<Item = u32>>(iter: I) -> Option<u32> {
let mut rng = crate::test::rng(411);
iter.choose_stable(&mut rng)
}
assert_eq!(choose([].iter().cloned()), None);
assert_eq!(choose(0..100), Some(40));
assert_eq!(choose(UnhintedIterator { iter: 0..100 }), Some(40));
assert_eq!(
choose(ChunkHintedIterator {
iter: 0..100,
chunk_size: 32,
chunk_remaining: 32,
hint_total_size: false,
}),
Some(40)
);
assert_eq!(
choose(ChunkHintedIterator {
iter: 0..100,
chunk_size: 32,
chunk_remaining: 32,
hint_total_size: true,
}),
Some(40)
);
assert_eq!(
choose(WindowHintedIterator {
iter: 0..100,
window_size: 32,
hint_total_size: false,
}),
Some(40)
);
assert_eq!(
choose(WindowHintedIterator {
iter: 0..100,
window_size: 32,
hint_total_size: true,
}),
Some(40)
);
}
#[test]
fn value_stability_choose_multiple() {
fn do_test<I: Iterator<Item = u32>>(iter: I, v: &[u32]) {
let mut rng = crate::test::rng(412);
let mut buf = [0u32; 8];
assert_eq!(iter.choose_multiple_fill(&mut rng, &mut buf), v.len());
assert_eq!(&buf[0..v.len()], v);
}
do_test(0..4, &[0, 1, 2, 3]);
do_test(0..8, &[0, 1, 2, 3, 4, 5, 6, 7]);
do_test(0..100, &[58, 78, 80, 92, 43, 8, 96, 7]);
#[cfg(feature = "alloc")]
{
fn do_test<I: Iterator<Item = u32>>(iter: I, v: &[u32]) {
let mut rng = crate::test::rng(412);
assert_eq!(iter.choose_multiple(&mut rng, v.len()), v);
}
do_test(0..4, &[0, 1, 2, 3]);
do_test(0..8, &[0, 1, 2, 3, 4, 5, 6, 7]);
do_test(0..100, &[58, 78, 80, 92, 43, 8, 96, 7]);
}
}
#[test]
#[cfg(feature = "std")]
fn test_multiple_weighted_edge_cases() {
use super::*;
let mut rng = crate::test::rng(413);
let choices = [('a', 2), ('b', 1), ('c', 0)];
for _ in 0..100 {
let result = choices
.choose_multiple_weighted(&mut rng, 2, |item| item.1)
.unwrap()
.collect::<Vec<_>>();
assert_eq!(result.len(), 2);
assert!(!result.iter().any(|val| val.0 == 'c'));
}
let choices = [('a', 0), ('b', 0), ('c', 0)];
let result = choices
.choose_multiple_weighted(&mut rng, 2, |item| item.1)
.unwrap()
.collect::<Vec<_>>();
assert_eq!(result.len(), 2);
let choices = [('a', -1), ('b', 1), ('c', 1)];
assert_eq!(
choices
.choose_multiple_weighted(&mut rng, 2, |item| item.1)
.unwrap_err(),
WeightedError::InvalidWeight
);
let choices = [];
let result = choices
.choose_multiple_weighted(&mut rng, 0, |_: &()| 0)
.unwrap()
.collect::<Vec<_>>();
assert_eq!(result.len(), 0);
let choices = [('a', core::f64::NAN), ('b', 1.0), ('c', 1.0)];
assert_eq!(
choices
.choose_multiple_weighted(&mut rng, 2, |item| item.1)
.unwrap_err(),
WeightedError::InvalidWeight
);
let choices = [('a', core::f64::INFINITY), ('b', 1.0), ('c', 1.0)];
for _ in 0..100 {
let result = choices
.choose_multiple_weighted(&mut rng, 2, |item| item.1)
.unwrap()
.collect::<Vec<_>>();
assert_eq!(result.len(), 2);
assert!(result.iter().any(|val| val.0 == 'a'));
}
let choices = [('a', core::f64::NEG_INFINITY), ('b', 1.0), ('c', 1.0)];
assert_eq!(
choices
.choose_multiple_weighted(&mut rng, 2, |item| item.1)
.unwrap_err(),
WeightedError::InvalidWeight
);
let choices = [('a', -0.0), ('b', 1.0), ('c', 1.0)];
assert!(choices
.choose_multiple_weighted(&mut rng, 2, |item| item.1)
.is_ok());
}
#[test]
#[cfg(feature = "std")]
fn test_multiple_weighted_distributions() {
use super::*;
let choices = [('a', 2), ('b', 1), ('c', 1)];
let mut rng = crate::test::rng(414);
let mut results = [0i32; 3];
let expected_results = [4167, 4167, 1666];
for _ in 0..10000 {
let result = choices
.choose_multiple_weighted(&mut rng, 2, |item| item.1)
.unwrap()
.collect::<Vec<_>>();
assert_eq!(result.len(), 2);
match (result[0].0, result[1].0) {
('a', 'b') | ('b', 'a') => {
results[0] += 1;
}
('a', 'c') | ('c', 'a') => {
results[1] += 1;
}
('b', 'c') | ('c', 'b') => {
results[2] += 1;
}
(_, _) => panic!("unexpected result"),
}
}
let mut diffs = results
.iter()
.zip(&expected_results)
.map(|(a, b)| (a - b).abs());
assert!(!diffs.any(|deviation| deviation > 100));
}
}