rand_functors/strategies/
population_sampler.rs

1use alloc::vec::Vec;
2
3use rand::distr::uniform::SampleUniform;
4use rand::distr::StandardUniform;
5use rand::prelude::*;
6
7use crate::{Inner, RandomStrategy, RandomVariable, RandomVariableRange};
8
9/// Produces a random subset (technically, submultiset) of possible outputs of
10/// the random process.
11#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
12pub struct PopulationSampler<const N: usize>;
13
14impl<const N: usize> PopulationSampler<N> {
15    #[inline(always)]
16    fn shrink_to_capacity<T: Inner>(mut f: Vec<T>, rng: &mut impl Rng) -> Vec<T> {
17        while f.len() > N {
18            let index = rng.random_range(0..f.len());
19            f.swap_remove(index);
20        }
21        f
22    }
23}
24
25impl<const N: usize> RandomStrategy for PopulationSampler<N> {
26    type Functor<I: Inner> = Vec<I>;
27
28    #[inline]
29    fn fmap<A: Inner, B: Inner, F: Fn(A) -> B>(f: Self::Functor<A>, func: F) -> Self::Functor<B> {
30        f.into_iter().map(func).collect()
31    }
32
33    #[inline]
34    fn fmap_rand<A: Inner, B: Inner, R: RandomVariable, F: Fn(A, R) -> B>(
35        f: Self::Functor<A>,
36        rng: &mut impl Rng,
37        func: F,
38    ) -> Self::Functor<B>
39    where
40        StandardUniform: Distribution<R>,
41    {
42        Self::shrink_to_capacity(
43            f.into_iter()
44                .flat_map(|a| R::sample_space().map(move |r| (a.clone(), r)))
45                .map(|(a, r)| func(a, r))
46                .collect(),
47            rng,
48        )
49    }
50
51    #[inline]
52    fn fmap_rand_range<A: Inner, B: Inner, R: RandomVariable + SampleUniform, F: Fn(A, R) -> B>(
53        f: Self::Functor<A>,
54        range: impl RandomVariableRange<R>,
55        rng: &mut impl Rng,
56        func: F,
57    ) -> Self::Functor<B>
58    where
59        StandardUniform: Distribution<R>,
60    {
61        Self::shrink_to_capacity(
62            f.into_iter()
63                .flat_map(|a| range.sample_space().map(move |r| (a.clone(), r)))
64                .map(|(a, r)| func(a, r))
65                .collect(),
66            rng,
67        )
68    }
69}