Skip to main content

checkito/
sample.rs

1use crate::{
2    SAMPLES,
3    generate::Generate,
4    shrink::{Shrink, Shrinkers},
5    state::{self, Modes, Sizes, State},
6};
7use core::iter;
8
9/// Configures the sampling process.
10///
11/// This struct is created by the [`Sample::sampler`] method and provides a
12/// builder-like interface for configuring how values are sampled.
13#[derive(Debug, Clone)]
14pub struct Sampler<G: ?Sized> {
15    /// The seed for the random number generator.
16    ///
17    /// Using the same seed will cause the sampler to produce the same sequence
18    /// of random values. It defaults to a random value.
19    pub seed: u64,
20    /// The range of sizes (`0.0..=1.0`) that will be gradually traversed while
21    /// generating values.
22    ///
23    /// Defaults to `0.0..=1.0`.
24    pub sizes: Sizes,
25    /// The number of samples to generate.
26    ///
27    /// Defaults to `SAMPLES`.
28    pub count: usize,
29    /// The generator that will provide the samples.
30    pub generator: G,
31}
32
33/// An iterator that yields random values from a generator.
34///
35/// This struct is created by the [`Sample::samples`] method.
36#[derive(Debug, Clone)]
37pub struct Samples<G: ?Sized>(Shrinkers<G>);
38
39/// An extension trait, implemented for all [`Generate`] types, that provides
40/// methods for sampling random values.
41pub trait Sample: Generate {
42    /// Creates a [`Sampler`] for this generator.
43    ///
44    /// The `Sampler` can be used to configure and control the sampling process.
45    fn sampler(self) -> Sampler<Self>
46    where
47        Self: Sized,
48    {
49        Sampler::new(self, state::seed())
50    }
51
52    /// Creates an iterator that generates `count` random values.
53    ///
54    /// The generated values will have progressively larger sizes. For more
55    /// control over the sampling process, see [`Sample::sampler`].
56    fn samples(self, count: usize) -> Samples<Self>
57    where
58        Self: Sized,
59    {
60        let mut sampler = self.sampler();
61        sampler.count = count;
62        sampler.samples()
63    }
64
65    /// Generates a single random value of a specific `size`.
66    ///
67    /// The `size` should be between `0.0` and `1.0`. For more control over the
68    /// sampling process, see [`Sample::sampler`].
69    fn sample(&self, size: f64) -> Self::Item {
70        self.sampler().sample(size)
71    }
72}
73
74impl<G: Generate + ?Sized> Sample for G {}
75
76impl<G> Sampler<G> {
77    pub(crate) const fn new(generator: G, seed: u64) -> Self {
78        Self {
79            generator,
80            seed,
81            sizes: Sizes::DEFAULT,
82            count: SAMPLES,
83        }
84    }
85}
86
87impl<G: Generate + ?Sized> Sampler<G> {
88    pub fn sample(&self, size: f64) -> G::Item {
89        let mut state = State::random(0, 1, size.into(), self.seed);
90        self.generator.generate(&mut state).item()
91    }
92}
93
94impl<G: Generate> Sampler<G> {
95    pub fn samples(self) -> Samples<G> {
96        let cardinality = self.generator.cardinality();
97        Samples::new(
98            self.generator,
99            Modes::with(self.count, self.sizes, self.seed, cardinality, Some(false)),
100        )
101    }
102}
103
104impl<G: Generate> Samples<G> {
105    pub(crate) fn new(generator: G, modes: Modes) -> Self {
106        Self(Shrinkers::new(generator, modes))
107    }
108}
109
110impl<G: Generate> Iterator for Samples<G> {
111    type Item = G::Item;
112
113    fn next(&mut self) -> Option<Self::Item> {
114        Some(self.0.next()?.item())
115    }
116
117    fn size_hint(&self) -> (usize, Option<usize>) {
118        self.0.size_hint()
119    }
120
121    fn count(self) -> usize {
122        self.0.count()
123    }
124
125    fn nth(&mut self, n: usize) -> Option<Self::Item> {
126        Some(self.0.nth(n)?.item())
127    }
128
129    fn last(self) -> Option<Self::Item> {
130        Some(self.0.last()?.item())
131    }
132}
133
134impl<G: Generate> DoubleEndedIterator for Samples<G> {
135    fn next_back(&mut self) -> Option<Self::Item> {
136        Some(self.0.next_back()?.item())
137    }
138
139    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
140        Some(self.0.nth_back(n)?.item())
141    }
142}
143
144impl<G: Generate> ExactSizeIterator for Samples<G> {
145    fn len(&self) -> usize {
146        self.0.len()
147    }
148}
149
150impl<G: Generate> iter::FusedIterator for Samples<G> {}