hegel/generators/
generators.rs

1use crate::test_case::{TestCase, labels};
2use ciborium::Value;
3use std::marker::PhantomData;
4use std::sync::Arc;
5
6/// A bundled schema + parse function for schema-based generation.
7///
8/// The lifetime `'a` ties the BasicGenerator to the generator that created it.
9/// `T: 'a` is required because the parse closure returns `T`.
10#[doc(hidden)]
11pub struct BasicGenerator<'a, T> {
12    schema: Value,
13    parse: Box<dyn Fn(Value) -> T + Send + Sync + 'a>,
14    _phantom: PhantomData<fn() -> T>,
15}
16
17impl<'a, T: 'a> BasicGenerator<'a, T> {
18    pub fn new<F: Fn(Value) -> T + Send + Sync + 'a>(schema: Value, f: F) -> Self {
19        BasicGenerator {
20            schema,
21            parse: Box::new(f),
22            _phantom: PhantomData,
23        }
24    }
25
26    pub fn schema(&self) -> &Value {
27        &self.schema
28    }
29
30    pub fn parse_raw(&self, raw: Value) -> T {
31        (self.parse)(raw)
32    }
33
34    /// Generate a value by sending the schema to the server and parsing the response.
35    pub fn do_draw(&self, tc: &TestCase) -> T {
36        self.parse_raw(super::generate_raw(tc, self.schema()))
37    }
38
39    /// Transform the output type by composing a function with the parse.
40    pub fn map<U: 'a, F: Fn(T) -> U + Send + Sync + 'a>(self, f: F) -> BasicGenerator<'a, U> {
41        let old_parse = self.parse;
42        BasicGenerator {
43            schema: self.schema,
44            parse: Box::new(move |raw| f(old_parse(raw))),
45            _phantom: PhantomData,
46        }
47    }
48}
49
50/// The core trait for all generators.
51///
52/// Generators produce values of type `T` and optionally provide a
53/// [`BasicGenerator`] for server-based generation via `as_basic()`.
54///
55/// Implementors override one or both of [`as_basic`](Self::as_basic) and
56/// [`do_draw`](Self::do_draw):
57/// - "Always basic" leaf generators override only `as_basic`. The default
58///   `do_draw` delegates to it.
59/// - Generators with a client-side fallback (filter, flat_map, lists with
60///   non-basic elements, …) override `do_draw`. They may also override
61///   `as_basic` to opt into schema generation when their inputs allow it.
62pub trait Generator<T>: Send + Sync {
63    /// Produce a value.
64    ///
65    /// The default delegates to [`as_basic`](Self::as_basic). Generators
66    /// without a basic form must override this.
67    #[doc(hidden)]
68    fn do_draw(&self, tc: &TestCase) -> T {
69        self.as_basic()
70            .expect("non-basic generator must override do_draw")
71            .do_draw(tc)
72    }
73
74    /// Return a BasicGenerator for schema-based generation, if possible.
75    #[doc(hidden)]
76    fn as_basic(&self) -> Option<BasicGenerator<'_, T>> {
77        None
78    }
79
80    /// Transform generated values using a function.
81    ///
82    /// When the source generator has a schema (i.e. `as_basic()` returns `Some`),
83    /// the schema is preserved and the function is composed into the parse step.
84    ///
85    /// # Example
86    ///
87    /// ```no_run
88    /// use hegel::generators::{self as gs, Generator};
89    ///
90    /// // Generate even integers by doubling
91    /// let evens = gs::integers::<i32>().map(|n| n * 2);
92    /// ```
93    fn map<U, F>(self, f: F) -> Mapped<T, U, F, Self>
94    where
95        Self: Sized,
96        F: Fn(T) -> U + Send + Sync,
97    {
98        Mapped {
99            source: self,
100            f: Arc::new(f),
101            _phantom: PhantomData,
102        }
103    }
104
105    /// Generate a value, then use it to choose or configure another generator.
106    ///
107    /// # Example
108    ///
109    /// ```no_run
110    /// use hegel::generators::{self as gs, Generator};
111    ///
112    /// // Generate a length, then a vec of exactly that length
113    /// let generator = gs::integers::<usize>()
114    ///     .min_value(1)
115    ///     .max_value(10)
116    ///     .flat_map(|len| gs::vecs(gs::integers::<i32>())
117    ///         .min_size(len)
118    ///         .max_size(len));
119    /// ```
120    fn flat_map<U, G, F>(self, f: F) -> FlatMapped<T, U, G, F, Self>
121    where
122        Self: Sized,
123        G: Generator<U>,
124        F: Fn(T) -> G + Send + Sync,
125    {
126        FlatMapped {
127            source: self,
128            f,
129            _phantom: PhantomData,
130        }
131    }
132
133    /// Only keep generated values that satisfy the predicate.
134    ///
135    /// Retries up to 3 times, then calls `assume(false)` to reject the test case.
136    ///
137    /// # Example
138    ///
139    /// ```no_run
140    /// use hegel::generators::{self as gs, Generator};
141    ///
142    /// // Generate integers, then filter out the even ones
143    /// let odds = gs::integers::<i32>().filter(|n| n % 2 != 0);
144    /// ```
145    fn filter<F>(self, predicate: F) -> Filtered<T, F, Self>
146    where
147        Self: Sized,
148        F: Fn(&T) -> bool + Send + Sync,
149    {
150        Filtered {
151            source: self,
152            predicate,
153            _phantom: PhantomData,
154        }
155    }
156
157    /// Return all possible values if this generator has a known finite value set.
158    ///
159    /// Used by [`Filtered`] as a fallback when random sampling fails: instead of
160    /// calling `assume(false)`, enumerate valid elements and pick one.
161    /// Mirrors Hypothesis's `SampledFromStrategy.do_filtered_draw` optimization.
162    #[doc(hidden)]
163    fn enumerate_values(&self) -> Option<Vec<T>> {
164        None
165    }
166
167    /// Convert this generator into a type-erased boxed generator.
168    ///
169    /// This is needed when you have generators of different concrete types
170    /// but the same output type and need to store them together, e.g. in a
171    /// `Vec` or when passing to [`one_of()`](super::one_of).
172    ///
173    /// # Example
174    ///
175    /// ```no_run
176    /// use hegel::generators::{self as gs, Generator};
177    ///
178    /// // Different generator types producing the same output type —
179    /// // boxing lets them be stored in a vec and passed to one_of
180    /// let generator = vec![
181    ///     gs::integers::<i32>().min_value(0).max_value(10).boxed(),
182    ///     gs::integers::<i32>().map(|n| n * 100).boxed(),
183    ///     gs::sampled_from(vec![1, 2, 3]).boxed(),
184    /// ];
185    /// ```
186    fn boxed<'a>(self) -> BoxedGenerator<'a, T>
187    where
188        Self: Sized + Send + Sync + 'a,
189    {
190        BoxedGenerator {
191            inner: Arc::new(self),
192        }
193    }
194}
195
196impl<T, G: Generator<T>> Generator<T> for &G {
197    fn do_draw(&self, tc: &TestCase) -> T {
198        (*self).do_draw(tc)
199    }
200
201    // nocov start
202    fn as_basic(&self) -> Option<BasicGenerator<'_, T>> {
203        (*self).as_basic()
204        // nocov end
205    }
206}
207
208/// Result of [`Generator::map`]. Preserves the schema when possible.
209pub struct Mapped<T, U, F, G> {
210    source: G,
211    f: Arc<F>,
212    _phantom: PhantomData<fn(T) -> U>,
213}
214
215impl<T, U, F, G> Generator<U> for Mapped<T, U, F, G>
216where
217    G: Generator<T>,
218    F: Fn(T) -> U + Send + Sync,
219{
220    fn do_draw(&self, tc: &TestCase) -> U {
221        if let Some(basic) = self.as_basic() {
222            basic.do_draw(tc)
223        } else {
224            tc.start_span(labels::MAPPED);
225            let result = (self.f)(self.source.do_draw(tc));
226            tc.stop_span(false);
227            result
228        }
229    }
230
231    fn as_basic(&self) -> Option<BasicGenerator<'_, U>> {
232        let source_basic = self.source.as_basic()?;
233        let f = Arc::clone(&self.f);
234        Some(source_basic.map(move |t| f(t)))
235    }
236
237    fn enumerate_values(&self) -> Option<Vec<U>> {
238        self.source
239            .enumerate_values()
240            .map(|vals| vals.into_iter().map(|v| (self.f)(v)).collect())
241    }
242}
243
244/// Result of [`Generator::flat_map`].
245pub struct FlatMapped<T, U, G2, F, G1> {
246    source: G1,
247    f: F,
248    _phantom: PhantomData<fn(T) -> (U, G2)>,
249}
250
251impl<T, U, G2, F, G1> Generator<U> for FlatMapped<T, U, G2, F, G1>
252where
253    G1: Generator<T>,
254    G2: Generator<U>,
255    F: Fn(T) -> G2 + Send + Sync,
256{
257    fn do_draw(&self, tc: &TestCase) -> U {
258        tc.start_span(labels::FLAT_MAP);
259        let intermediate = self.source.do_draw(tc);
260        let next_gen = (self.f)(intermediate);
261        let result = next_gen.do_draw(tc);
262        tc.stop_span(false);
263        result
264    }
265}
266
267/// Result of [`Generator::filter`].
268pub struct Filtered<T, F, G> {
269    source: G,
270    predicate: F,
271    _phantom: PhantomData<fn() -> T>,
272}
273
274impl<T, F, G> Generator<T> for Filtered<T, F, G>
275where
276    T: Clone + Send + Sync,
277    G: Generator<T>,
278    F: Fn(&T) -> bool + Send + Sync,
279{
280    fn do_draw(&self, tc: &TestCase) -> T {
281        if let Some(basic) = self.as_basic() {
282            return basic.do_draw(tc);
283        }
284        for _ in 0..3 {
285            tc.start_span(labels::FILTER);
286            let value = self.source.do_draw(tc);
287            if (self.predicate)(&value) {
288                tc.stop_span(false);
289                return value;
290            }
291            tc.stop_span(true);
292        }
293        if self
294            .enumerate_values()
295            .is_some_and(|valid| valid.is_empty())
296        {
297            panic!(
298                "Unsatisfiable filter: all values from the source generator \
299                 are rejected by the filter predicate"
300            );
301        }
302        tc.assume(false);
303        unreachable!()
304    }
305
306    fn as_basic(&self) -> Option<BasicGenerator<'_, T>> {
307        let valid = self.enumerate_values()?;
308        if valid.is_empty() {
309            return None;
310        }
311        use crate::cbor_utils::cbor_map;
312        let schema = cbor_map! {
313            "type" => "integer",
314            "min_value" => 0u64,
315            "max_value" => (valid.len() - 1) as u64
316        };
317        Some(BasicGenerator::new(schema, move |raw| {
318            let index: usize = super::deserialize_value(raw);
319            valid[index].clone()
320        }))
321    }
322
323    fn enumerate_values(&self) -> Option<Vec<T>> {
324        self.source
325            .enumerate_values()
326            .map(|vals| vals.into_iter().filter(|v| (self.predicate)(v)).collect())
327    }
328}
329
330/// A type-erased generator with a lifetime parameter.
331pub struct BoxedGenerator<'a, T> {
332    pub(super) inner: Arc<dyn Generator<T> + Send + Sync + 'a>,
333}
334
335impl<T> Clone for BoxedGenerator<'_, T> {
336    fn clone(&self) -> Self {
337        BoxedGenerator {
338            inner: Arc::clone(&self.inner),
339        }
340    }
341}
342
343impl<T> Generator<T> for BoxedGenerator<'_, T> {
344    fn do_draw(&self, tc: &TestCase) -> T {
345        self.inner.do_draw(tc)
346    }
347
348    fn as_basic(&self) -> Option<BasicGenerator<'_, T>> {
349        self.inner.as_basic()
350    }
351
352    fn enumerate_values(&self) -> Option<Vec<T>> {
353        self.inner.enumerate_values()
354    }
355
356    fn boxed<'b>(self) -> BoxedGenerator<'b, T>
357    where
358        Self: Sized + Send + Sync + 'b,
359    {
360        BoxedGenerator { inner: self.inner }
361    }
362}
hegel/generators/generators.rs

hegel/generators/
generators.rs
