rand_functors/strategies/
counter.rs

1use std::collections::hash_map::RandomState;
2use std::collections::HashMap;
3use std::hash::BuildHasher;
4use std::marker::PhantomData;
5
6use num::traits::{NumAssign, Unsigned};
7use num::Integer;
8use rand::distr::uniform::SampleUniform;
9use rand::distr::StandardUniform;
10use rand::prelude::*;
11
12use crate::{
13    FlattenableRandomStrategy, Inner, RandomStrategy, RandomVariable, RandomVariableRange,
14};
15
16/// Produces all possible outputs of the random process, with repetition, stored
17/// in a [`HashMap`].
18///
19/// `Counter` is optimal in scenarios where certain operations will map many
20/// inputs to the same output. Examples include conditionally zeroing out a
21/// field of a struct or the use of functions like `saturating_add` or
22/// `saturating_mul`.
23#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
24pub struct Counter<
25    S: BuildHasher + Default = RandomState,
26    N: Clone + Default + NumAssign + Unsigned = usize,
27> {
28    count_phantom: PhantomData<N>,
29    hasher_phantom: PhantomData<S>,
30}
31
32impl<S: BuildHasher + Default, N: Clone + Default + NumAssign + Unsigned> RandomStrategy
33    for Counter<S, N>
34{
35    type Functor<I: Inner> = HashMap<I, N, S>;
36
37    #[inline]
38    fn fmap<A: Inner, B: Inner, F: Fn(A) -> B>(f: Self::Functor<A>, func: F) -> Self::Functor<B> {
39        // Constructing a new HashMap is necessary, as there may be fewer new
40        // keys than old keys, which requires merging some or all counts.
41        let mut new_functor = Self::Functor::with_capacity_and_hasher(f.len(), Default::default());
42        f.into_iter()
43            .map(|(i, count)| (func(i), count))
44            .for_each(|(o, count)| {
45                *new_functor.entry(o).or_insert(N::zero()) += count;
46            });
47        new_functor
48    }
49
50    #[inline]
51    fn fmap_rand<A: Inner, B: Inner, R: RandomVariable, F: Fn(A, R) -> B>(
52        f: Self::Functor<A>,
53        _: &mut impl Rng,
54        func: F,
55    ) -> Self::Functor<B>
56    where
57        StandardUniform: Distribution<R>,
58    {
59        let mut new_functor = Self::Functor::with_capacity_and_hasher(f.len(), Default::default());
60        f.into_iter()
61            .flat_map(|a| R::sample_space().map(move |r| (a.clone(), r)))
62            .map(|((a, c), r)| (func(a, r), c))
63            .for_each(|(b, count)| {
64                *new_functor.entry(b).or_insert(N::zero()) += count;
65            });
66        new_functor
67    }
68
69    #[inline]
70    fn fmap_rand_range<A: Inner, B: Inner, R: RandomVariable + SampleUniform, F: Fn(A, R) -> B>(
71        f: Self::Functor<A>,
72        range: impl RandomVariableRange<R>,
73        _: &mut impl Rng,
74        func: F,
75    ) -> Self::Functor<B>
76    where
77        StandardUniform: Distribution<R>,
78    {
79        let mut new_functor = Self::Functor::with_capacity_and_hasher(f.len(), Default::default());
80        f.into_iter()
81            .flat_map(|a| range.sample_space().map(move |r| (a.clone(), r)))
82            .map(|((a, c), r)| (func(a, r), c))
83            .for_each(|(b, count)| {
84                *new_functor.entry(b).or_insert(N::zero()) += count;
85            });
86        new_functor
87    }
88}
89
90impl<S: BuildHasher + Default, N: Clone + Default + Integer + NumAssign + Unsigned>
91    FlattenableRandomStrategy for Counter<S, N>
92{
93    #[inline]
94    fn fmap_flat<A: Inner, B: Inner, F: FnMut(A) -> Self::Functor<B>>(
95        f: Self::Functor<A>,
96        mut func: F,
97    ) -> Self::Functor<B> {
98        let mut new_functor = Self::Functor::with_capacity_and_hasher(f.len(), Default::default());
99        let children = f
100            .into_iter()
101            .map(|(i, outer_count)| {
102                let functor = func(i);
103                let inner_count_sum = functor
104                    .values()
105                    .fold(N::zero(), |sum, inner_count| sum + inner_count.clone());
106                (functor, outer_count, inner_count_sum)
107            })
108            .collect::<Vec<_>>();
109        let Some(inner_count_lcm) =
110            children
111                .iter()
112                .fold(None, |lcm: Option<N>, (_, _, inner_count_sum)| {
113                    if let Some(lcm) = lcm {
114                        Some(num::integer::lcm(lcm, inner_count_sum.clone()))
115                    } else {
116                        Some(inner_count_sum.clone())
117                    }
118                })
119        else {
120            return new_functor;
121        };
122        for (child, outer_count, inner_count_sum) in children {
123            let scaling = inner_count_lcm.clone() / inner_count_sum;
124            for (output, inner_count) in child {
125                *new_functor.entry(output).or_insert(N::zero()) +=
126                    scaling.clone() * inner_count * outer_count.clone();
127            }
128        }
129        new_functor
130    }
131}
rand_functors/strategies/counter.rs

rand_functors/strategies/
counter.rs
