hegel/generators/
collections.rs

1use super::{BasicGenerator, BoxedGenerator, Collection, Generator, TestCase, labels};
2use crate::cbor_utils::{cbor_map, map_insert};
3use ciborium::Value;
4use std::collections::{HashMap, HashSet};
5use std::hash::Hash;
6use std::marker::PhantomData;
7use std::sync::Arc;
8
9/// Generator for `Vec<T>`. Created by [`vecs()`].
10pub struct VecGenerator<G, T> {
11    pub(crate) elements: G,
12    pub(crate) min_size: usize,
13    pub(crate) max_size: Option<usize>,
14    pub(crate) unique_by: Option<fn(&T, &T) -> bool>,
15    pub(crate) _phantom: PhantomData<fn(T)>,
16}
17
18impl<G, T> VecGenerator<G, T> {
19    /// Set the minimum number of elements.
20    pub fn min_size(mut self, min_size: usize) -> Self {
21        self.min_size = min_size;
22        self
23    }
24
25    /// Set the maximum number of elements.
26    pub fn max_size(mut self, max_size: usize) -> Self {
27        self.max_size = Some(max_size);
28        self
29    }
30}
31
32impl<G, T: PartialEq> VecGenerator<G, T> {
33    /// Require all elements to be unique.
34    pub fn unique(mut self, unique: bool) -> Self {
35        self.unique_by = if unique {
36            Some(<T as PartialEq>::eq)
37        } else {
38            None
39        };
40        self
41    }
42}
43
44impl<T, G> Generator<Vec<T>> for VecGenerator<G, T>
45where
46    G: Generator<T>,
47{
48    fn do_draw(&self, tc: &TestCase) -> Vec<T> {
49        if let Some(max) = self.max_size {
50            assert!(self.min_size <= max, "Cannot have max_size < min_size");
51        }
52        if let Some(basic) = self.as_basic() {
53            basic.do_draw(tc)
54        } else {
55            tc.start_span(labels::LIST);
56            let mut collection = Collection::new(tc, self.min_size, self.max_size);
57            let mut result = Vec::new();
58            while collection.more() {
59                let element = self.elements.do_draw(tc);
60                if let Some(eq_fn) = &self.unique_by {
61                    if result.iter().any(|existing| eq_fn(existing, &element)) {
62                        collection.reject(Some("duplicate element"));
63                        continue;
64                    }
65                }
66                result.push(element);
67            }
68            tc.stop_span(false);
69            result
70        }
71    }
72
73    fn as_basic(&self) -> Option<BasicGenerator<'_, Vec<T>>> {
74        if let Some(max) = self.max_size {
75            assert!(self.min_size <= max, "Cannot have max_size < min_size");
76        }
77        if self.unique_by.is_some() {
78            return None;
79        }
80        let basic = self.elements.as_basic()?;
81
82        let mut schema = cbor_map! {
83            "type" => "list",
84            "unique" => false,
85            "elements" => basic.schema().clone(),
86            "min_size" => self.min_size as u64
87        };
88
89        if let Some(max) = self.max_size {
90            map_insert(&mut schema, "max_size", max as u64);
91        }
92
93        Some(BasicGenerator::new(schema, move |raw| {
94            let Value::Array(arr) = raw else {
95                panic!("Expected array, got {:?}", raw) // nocov
96            };
97            arr.into_iter().map(|v| basic.parse_raw(v)).collect()
98        }))
99    }
100}
101
102/// Generate vectors with elements from the given generator.
103///
104/// See [`VecGenerator`] for builder methods.
105///
106/// # Example
107///
108/// ```no_run
109/// use hegel::generators as gs;
110///
111/// #[hegel::test]
112/// fn my_test(tc: hegel::TestCase) {
113///     let v: Vec<i32> = tc.draw(gs::vecs(gs::integers())
114///         .min_size(1)
115///         .max_size(10));
116///     assert!(!v.is_empty() && v.len() <= 10);
117/// }
118/// ```
119pub fn vecs<T, G: Generator<T>>(elements: G) -> VecGenerator<G, T> {
120    VecGenerator {
121        elements,
122        min_size: 0,
123        max_size: None,
124        unique_by: None,
125        _phantom: PhantomData,
126    }
127}
128
129/// Generator for `HashSet<T>`. Created by [`hashsets()`].
130pub struct HashSetGenerator<G, T> {
131    elements: G,
132    min_size: usize,
133    max_size: Option<usize>,
134    _phantom: PhantomData<fn(T)>,
135}
136
137impl<G, T> HashSetGenerator<G, T> {
138    /// Set the minimum number of elements.
139    pub fn min_size(mut self, min_size: usize) -> Self {
140        self.min_size = min_size;
141        self
142    }
143
144    /// Set the maximum number of elements.
145    pub fn max_size(mut self, max_size: usize) -> Self {
146        self.max_size = Some(max_size);
147        self
148    }
149}
150
151impl<T, G> Generator<HashSet<T>> for HashSetGenerator<G, T>
152where
153    G: Generator<T>,
154    T: Eq + Hash,
155{
156    fn do_draw(&self, tc: &TestCase) -> HashSet<T> {
157        if let Some(max) = self.max_size {
158            assert!(self.min_size <= max, "Cannot have max_size < min_size");
159        }
160        if let Some(basic) = self.as_basic() {
161            basic.do_draw(tc)
162        } else {
163            // nocov start
164            tc.start_span(labels::SET);
165            let mut collection = Collection::new(tc, self.min_size, self.max_size);
166            let mut set = HashSet::new();
167            while collection.more() {
168                let element = self.elements.do_draw(tc);
169                if !set.insert(element) {
170                    collection.reject(Some("duplicate element"));
171                    // nocov end
172                }
173            }
174            // nocov start
175            assert!(set.len() >= self.min_size);
176            tc.stop_span(false);
177            set
178            // nocov end
179        }
180    }
181
182    fn as_basic(&self) -> Option<BasicGenerator<'_, HashSet<T>>> {
183        if let Some(max) = self.max_size {
184            assert!(self.min_size <= max, "Cannot have max_size < min_size");
185        }
186        let basic = self.elements.as_basic()?;
187
188        let mut schema = cbor_map! {
189            "type" => "list",
190            "unique" => true,
191            "elements" =>  basic.schema().clone(),
192            "min_size" => self.min_size as u64
193        };
194
195        if let Some(max) = self.max_size {
196            map_insert(&mut schema, "max_size", max as u64);
197        }
198
199        Some(BasicGenerator::new(schema, move |raw| {
200            let Value::Array(arr) = raw else {
201                panic!("Expected array, got {:?}", raw) // nocov
202            };
203            arr.into_iter().map(|v| basic.parse_raw(v)).collect()
204        }))
205    }
206}
207
208/// Generate hash sets with elements from the given generator.
209///
210/// See [`HashSetGenerator`] for builder methods.
211pub fn hashsets<T, G: Generator<T>>(elements: G) -> HashSetGenerator<G, T> {
212    HashSetGenerator {
213        elements,
214        min_size: 0,
215        max_size: None,
216        _phantom: PhantomData,
217    }
218}
219
220/// Generator for `HashMap<K, V>`. Created by [`hashmaps()`].
221pub struct HashMapGenerator<K, V, KT, VT> {
222    keys: K,
223    values: V,
224    min_size: usize,
225    max_size: Option<usize>,
226    _phantom: PhantomData<fn(KT, VT)>,
227}
228
229impl<K, V, KT, VT> HashMapGenerator<K, V, KT, VT> {
230    /// Set the minimum number of entries.
231    pub fn min_size(mut self, min_size: usize) -> Self {
232        self.min_size = min_size;
233        self
234    }
235
236    /// Set the maximum number of entries.
237    pub fn max_size(mut self, max_size: usize) -> Self {
238        self.max_size = Some(max_size);
239        self
240    }
241}
242
243impl<K, V, KT, VT> Generator<HashMap<KT, VT>> for HashMapGenerator<K, V, KT, VT>
244where
245    K: Generator<KT>,
246    V: Generator<VT>,
247    KT: Eq + std::hash::Hash,
248{
249    fn do_draw(&self, tc: &TestCase) -> HashMap<KT, VT> {
250        if let Some(max) = self.max_size {
251            assert!(self.min_size <= max, "Cannot have max_size < min_size");
252        }
253        if let Some(basic) = self.as_basic() {
254            basic.do_draw(tc)
255        } else {
256            // nocov start
257            tc.start_span(labels::MAP);
258            let mut collection = Collection::new(tc, self.min_size, self.max_size);
259            let mut map = HashMap::new();
260            while collection.more() {
261                let key = self.keys.do_draw(tc);
262                match map.entry(key) {
263                    std::collections::hash_map::Entry::Occupied(_) => {
264                        collection.reject(Some("duplicate key"));
265                    }
266                    std::collections::hash_map::Entry::Vacant(entry) => {
267                        let value = self.values.do_draw(tc);
268                        entry.insert(value);
269                        // nocov end
270                    }
271                }
272            }
273            // nocov start
274            assert!(map.len() >= self.min_size);
275            tc.stop_span(false);
276            map
277            // nocov end
278        }
279    }
280
281    fn as_basic(&self) -> Option<BasicGenerator<'_, HashMap<KT, VT>>> {
282        if let Some(max) = self.max_size {
283            assert!(self.min_size <= max, "Cannot have max_size < min_size");
284        }
285        let keys_basic = self.keys.as_basic()?;
286        let values_basic = self.values.as_basic()?;
287
288        let mut schema = cbor_map! {
289            "type" => "dict",
290            "keys" => keys_basic.schema().clone(),
291            "values" => values_basic.schema().clone(),
292            "min_size" => self.min_size as u64
293        };
294
295        if let Some(max) = self.max_size {
296            map_insert(&mut schema, "max_size", max as u64);
297        }
298
299        Some(BasicGenerator::new(schema, move |raw| {
300            // schema expects format [[key, value], ...]
301            let values = match raw {
302                Value::Array(arr) => arr,
303                _ => panic!("Expected array, got {:?}", raw), // nocov
304            };
305
306            let mut map = HashMap::new();
307            for value_raw in values {
308                let value = match value_raw {
309                    Value::Array(arr) => arr,
310                    _ => panic!("Expected array, got {:?}", value_raw), // nocov
311                };
312                let mut iter = value.into_iter();
313                let raw_k = iter.next().unwrap();
314                let raw_v = iter.next().unwrap();
315
316                let key = keys_basic.parse_raw(raw_k);
317                let value = values_basic.parse_raw(raw_v);
318
319                map.insert(key, value);
320            }
321            map
322        }))
323    }
324}
325
326/// Generate hash maps.
327///
328/// See [`HashMapGenerator`] for builder methods.
329///
330/// # Example
331///
332/// ```ignore
333/// use hegel::generators as gs;
334/// use std::collections::HashMap;
335///
336/// let map: HashMap<i32, String> = tc.draw(gs::hashmaps(gs::integers(), gs::text()));
337/// ```
338pub fn hashmaps<KT, VT, K: Generator<KT>, V: Generator<VT>>(
339    keys: K,
340    values: V,
341) -> HashMapGenerator<K, V, KT, VT> {
342    HashMapGenerator {
343        keys,
344        values,
345        min_size: 0,
346        max_size: None,
347        _phantom: PhantomData,
348    }
349}
350
351pub(crate) struct MappedToValue<T, G> {
352    inner: G,
353    _phantom: PhantomData<fn() -> T>,
354}
355
356impl<T: serde::Serialize, G: Generator<T>> Generator<Value> for MappedToValue<T, G> {
357    // nocov start
358    fn do_draw(&self, tc: &TestCase) -> Value {
359        crate::cbor_utils::cbor_serialize(&self.inner.do_draw(tc))
360        // nocov end
361    }
362
363    // nocov start
364    fn as_basic(&self) -> Option<BasicGenerator<'_, Value>> {
365        let inner_basic = self.inner.as_basic()?;
366        let schema = inner_basic.schema().clone();
367        Some(BasicGenerator::new(schema, move |raw| {
368            let t_val = inner_basic.parse_raw(raw);
369            crate::cbor_utils::cbor_serialize(&t_val)
370            // nocov end
371        }))
372    }
373}
374
375/// Builder for fixed-key dictionary generators. Created by [`fixed_dicts()`].
376///
377/// Add fields with [`field()`](Self::field), then call [`build()`](Self::build)
378/// to get the generator.
379pub struct FixedDictBuilder<'a> {
380    fields: Vec<(String, BoxedGenerator<'a, Value>)>,
381}
382
383impl<'a> FixedDictBuilder<'a> {
384    /// Add a field with a name and generator.
385    // nocov start
386    pub fn field<T, G>(mut self, name: &str, generator: G) -> Self
387    where
388        G: Generator<T> + Send + Sync + 'a,
389        T: serde::Serialize + 'a,
390        // nocov end
391    {
392        // nocov start
393        let boxed = BoxedGenerator {
394            inner: Arc::new(MappedToValue {
395                inner: generator,
396                _phantom: PhantomData,
397                // nocov end
398            }),
399        };
400        self.fields.push((name.to_string(), boxed)); // nocov
401        self // nocov
402    }
403
404    /// Build the generator.
405    // nocov start
406    pub fn build(self) -> FixedDictGenerator<'a> {
407        FixedDictGenerator {
408            fields: self.fields,
409            // nocov end
410        }
411    }
412}
413
414/// Generator for CBOR maps with fixed keys. Created via [`FixedDictBuilder`].
415pub struct FixedDictGenerator<'a> {
416    fields: Vec<(String, BoxedGenerator<'a, Value>)>,
417}
418
419impl Generator<Value> for FixedDictGenerator<'_> {
420    // nocov start
421    fn do_draw(&self, tc: &TestCase) -> Value {
422        if let Some(basic) = self.as_basic() {
423            basic.do_draw(tc)
424        } else {
425            tc.start_span(labels::FIXED_DICT);
426            let entries: Vec<(Value, Value)> = self
427                .fields
428                .iter()
429                .map(|(name, g)| (Value::Text(name.clone()), g.do_draw(tc)))
430                .collect();
431            tc.stop_span(false);
432            Value::Map(entries)
433            // nocov end
434        }
435    }
436
437    // nocov start
438    fn as_basic(&self) -> Option<BasicGenerator<'_, Value>> {
439        let basics: Vec<BasicGenerator<'_, Value>> = self
440            .fields
441            .iter()
442            .map(|(_, g)| g.as_basic())
443            .collect::<Option<Vec<_>>>()?;
444
445        let schemas: Vec<Value> = basics.iter().map(|b| b.schema().clone()).collect();
446
447        let schema = cbor_map! {
448            "type" => "tuple",
449            "elements" => Value::Array(schemas)
450        // nocov end
451        };
452
453        let field_names: Vec<String> = self.fields.iter().map(|(name, _)| name.clone()).collect(); // nocov
454
455        // nocov start
456        Some(BasicGenerator::new(schema, move |raw| {
457            let arr = match raw {
458                Value::Array(arr) => arr,
459                _ => panic!("Expected array from tuple schema, got {:?}", raw),
460                // nocov end
461            };
462
463            // nocov start
464            let entries: Vec<(Value, Value)> = field_names
465                .iter()
466                .zip(basics.iter())
467                .zip(arr)
468                .map(|((name, basic), val)| (Value::Text(name.clone()), basic.parse_raw(val)))
469                .collect();
470            Value::Map(entries)
471            // nocov end
472        }))
473    }
474}
475
476/// Create a generator for dictionaries with fixed keys.
477///
478/// See [`FixedDictBuilder`] for builder methods.
479///
480/// # Example
481///
482/// ```no_run
483/// use hegel::generators::{self as gs, Generator};
484///
485/// let generator = gs::fixed_dicts()
486///     .field("name", gs::text())
487///     .field("age", gs::integers::<u32>())
488///     .build();
489/// ```
490// nocov start
491pub fn fixed_dicts<'a>() -> FixedDictBuilder<'a> {
492    FixedDictBuilder { fields: Vec::new() }
493    // nocov end
494}
495
496/// Generator for fixed-size arrays `[T; N]`. Created by [`arrays()`].
497pub struct ArrayGenerator<G, T, const N: usize> {
498    element: G,
499    _phantom: PhantomData<fn() -> T>,
500}
501
502impl<G, T, const N: usize> ArrayGenerator<G, T, N> {
503    #[doc(hidden)]
504    pub fn new(element: G) -> Self {
505        ArrayGenerator {
506            element,
507            _phantom: PhantomData,
508        }
509    }
510}
511
512/// Generate fixed-size arrays `[T; N]` with elements from the given generator.
513pub fn arrays<G: Generator<T> + Send + Sync, T, const N: usize>(
514    element: G,
515) -> ArrayGenerator<G, T, N> {
516    ArrayGenerator::new(element)
517}
518
519impl<G: Generator<T> + Send + Sync, T, const N: usize> Generator<[T; N]>
520    for ArrayGenerator<G, T, N>
521{
522    fn do_draw(&self, tc: &TestCase) -> [T; N] {
523        if let Some(basic) = self.as_basic() {
524            basic.do_draw(tc)
525        } else {
526            tc.start_span(labels::TUPLE);
527            let result = std::array::from_fn(|_| self.element.do_draw(tc));
528            tc.stop_span(false);
529            result
530        }
531    }
532
533    fn as_basic(&self) -> Option<BasicGenerator<'_, [T; N]>> {
534        let basic = self.element.as_basic()?;
535
536        let elements = Value::Array((0..N).map(|_| basic.schema().clone()).collect());
537        let schema = cbor_map! {
538            "type" => "tuple",
539            "elements" => elements
540        };
541
542        Some(BasicGenerator::new(schema, move |raw| {
543            let arr = match raw {
544                Value::Array(arr) => arr,
545                _ => panic!("Expected array from tuple schema, got {:?}", raw), // nocov
546            };
547            assert_eq!(arr.len(), N);
548            let mut iter = arr.into_iter();
549            std::array::from_fn(|_| basic.parse_raw(iter.next().unwrap()))
550        }))
551    }
552}
hegel/generators/collections.rs

hegel/generators/
collections.rs
