hegel/generators/
collections.rs

1use super::{Collection, Generator, TestCase, labels};
2use crate::control::hegel_internal_assert;
3use crate::test_case::invalid_argument;
4use std::collections::{HashMap, HashSet};
5use std::hash::Hash;
6use std::marker::PhantomData;
7
8/// Generator for `Vec<T>`. Created by [`vecs()`].
9pub struct VecGenerator<G, T> {
10    pub(crate) elements: G,
11    pub(crate) min_size: usize,
12    pub(crate) max_size: Option<usize>,
13    pub(crate) unique_by: Option<fn(&T, &T) -> bool>,
14    pub(crate) _phantom: PhantomData<fn(T)>,
15}
16
17impl<G, T> VecGenerator<G, T> {
18    /// Set the minimum number of elements.
19    pub fn min_size(mut self, min_size: usize) -> Self {
20        self.min_size = min_size;
21        self
22    }
23
24    /// Set the maximum number of elements.
25    pub fn max_size(mut self, max_size: usize) -> Self {
26        self.max_size = Some(max_size);
27        self
28    }
29}
30
31impl<G, T: PartialEq> VecGenerator<G, T> {
32    /// Require all elements to be unique.
33    pub fn unique(mut self, unique: bool) -> Self {
34        self.unique_by = if unique {
35            Some(<T as PartialEq>::eq)
36        } else {
37            None
38        };
39        self
40    }
41}
42
43impl<T, G> Generator<Vec<T>> for VecGenerator<G, T>
44where
45    G: Generator<T>,
46{
47    fn do_draw(&self, tc: &TestCase) -> Vec<T> {
48        if let Some(max) = self.max_size {
49            if self.min_size > max {
50                invalid_argument!("Cannot have max_size < min_size");
51            }
52        }
53        tc.start_span(labels::LIST);
54        let mut collection = Collection::new(tc, self.min_size, self.max_size);
55        let mut result = Vec::new();
56        while collection.more() {
57            let element = self.elements.do_draw(tc);
58            if let Some(eq_fn) = &self.unique_by {
59                if result.iter().any(|existing| eq_fn(existing, &element)) {
60                    collection.reject(Some("duplicate element"));
61                    continue;
62                }
63            }
64            result.push(element);
65        }
66        tc.stop_span(false);
67        result
68    }
69}
70
71/// Generate vectors with elements from the given generator.
72///
73/// See [`VecGenerator`] for builder methods.
74///
75/// # Example
76///
77/// ```no_run
78/// use hegel::generators as gs;
79///
80/// #[hegel::test]
81/// fn my_test(tc: hegel::TestCase) {
82///     let v: Vec<i32> = tc.draw(gs::vecs(gs::integers())
83///         .min_size(1)
84///         .max_size(10));
85///     assert!(!v.is_empty() && v.len() <= 10);
86/// }
87/// ```
88pub fn vecs<T, G: Generator<T>>(elements: G) -> VecGenerator<G, T> {
89    VecGenerator {
90        elements,
91        min_size: 0,
92        max_size: None,
93        unique_by: None,
94        _phantom: PhantomData,
95    }
96}
97
98/// Generator for `HashSet<T>`. Created by [`hashsets()`].
99pub struct HashSetGenerator<G, T> {
100    elements: G,
101    min_size: usize,
102    max_size: Option<usize>,
103    _phantom: PhantomData<fn(T)>,
104}
105
106impl<G, T> HashSetGenerator<G, T> {
107    /// Set the minimum number of elements.
108    pub fn min_size(mut self, min_size: usize) -> Self {
109        self.min_size = min_size;
110        self
111    }
112
113    /// Set the maximum number of elements.
114    pub fn max_size(mut self, max_size: usize) -> Self {
115        self.max_size = Some(max_size);
116        self
117    }
118}
119
120/// The largest enumerated value pool [`HashSetGenerator`] will draw
121/// without replacement from. Mirrors the bound the engine's old
122/// unique-sampled-list strategy used.
123const MAX_UNIQUE_POOL: usize = 10_000;
124
125impl<T, G> Generator<HashSet<T>> for HashSetGenerator<G, T>
126where
127    G: Generator<T>,
128    T: Eq + Hash,
129{
130    fn do_draw(&self, tc: &TestCase) -> HashSet<T> {
131        if let Some(max) = self.max_size {
132            if self.min_size > max {
133                invalid_argument!("Cannot have max_size < min_size");
134            }
135        }
136        tc.start_span(labels::SET);
137        let set = match self.enumerated_pool() {
138            Some(pool) => self.draw_from_pool(tc, pool),
139            None => self.draw_by_rejection(tc),
140        };
141        tc.stop_span(false);
142        set
143    }
144}
145
146impl<T, G> HashSetGenerator<G, T>
147where
148    G: Generator<T>,
149    T: Eq + Hash,
150{
151    /// The distinct values of an enumerable element generator, in first
152    /// occurrence order (which the shrinker treats as simplest-first), when
153    /// there are few enough of them to draw without replacement.
154    fn enumerated_pool(&self) -> Option<Vec<T>> {
155        let values = self.elements.enumerate_values()?;
156        if values.is_empty() || values.len() > MAX_UNIQUE_POOL {
157            return None;
158        }
159        let mut by_hash: HashMap<u64, Vec<usize>> = HashMap::new();
160        let mut pool: Vec<T> = Vec::new();
161        for v in values {
162            let bucket = by_hash.entry(fingerprint(&v)).or_default();
163            if bucket.iter().any(|&i| pool[i] == v) {
164                continue;
165            }
166            bucket.push(pool.len());
167            pool.push(v);
168        }
169        Some(pool)
170    }
171
172    /// Draw set elements as indices into a shrinking pool of the remaining
173    /// values, avoiding the coupon-collector problem when the set must
174    /// contain most of a small alphabet. Port of Hypothesis's
175    /// `UniqueSampledListStrategy`.
176    fn draw_from_pool(&self, tc: &TestCase, mut remaining: Vec<T>) -> HashSet<T> {
177        if self.min_size > remaining.len() {
178            invalid_argument!(
179                "Cannot generate a set: min_size {} is larger than the {} distinct values the element generator can produce",
180                self.min_size,
181                remaining.len()
182            );
183        }
184        let effective_max = self
185            .max_size
186            .map_or(remaining.len(), |m| m.min(remaining.len()));
187        let mut collection = Collection::new(tc, self.min_size, Some(effective_max));
188        let mut set = HashSet::new();
189        loop {
190            if remaining.is_empty() || !collection.more() {
191                break;
192            }
193            let j = tc.generate_integer_i64(0, remaining.len() as i64 - 1) as usize;
194            set.insert(remaining.remove(j));
195        }
196        set
197    }
198
199    fn draw_by_rejection(&self, tc: &TestCase) -> HashSet<T> {
200        let mut collection = Collection::new(tc, self.min_size, self.max_size);
201        let mut set = HashSet::new();
202        while collection.more() {
203            let element = self.elements.do_draw(tc);
204            if !set.insert(element) {
205                collection.reject(Some("duplicate element"));
206            }
207        }
208        hegel_internal_assert!(set.len() >= self.min_size);
209        set
210    }
211}
212
213/// A hashable stand-in for a value that is only `Eq + Hash`, used to dedup
214/// the enumerated pool.
215fn fingerprint<T: Eq + Hash>(v: &T) -> u64 {
216    use std::hash::{DefaultHasher, Hasher};
217    let mut h = DefaultHasher::new();
218    v.hash(&mut h);
219    h.finish()
220}
221
222/// Generate hash sets with elements from the given generator.
223///
224/// See [`HashSetGenerator`] for builder methods.
225pub fn hashsets<T, G: Generator<T>>(elements: G) -> HashSetGenerator<G, T> {
226    HashSetGenerator {
227        elements,
228        min_size: 0,
229        max_size: None,
230        _phantom: PhantomData,
231    }
232}
233
234/// Generator for `HashMap<K, V>`. Created by [`hashmaps()`].
235pub struct HashMapGenerator<K, V, KT, VT> {
236    keys: K,
237    values: V,
238    min_size: usize,
239    max_size: Option<usize>,
240    _phantom: PhantomData<fn(KT, VT)>,
241}
242
243impl<K, V, KT, VT> HashMapGenerator<K, V, KT, VT> {
244    /// Set the minimum number of entries.
245    pub fn min_size(mut self, min_size: usize) -> Self {
246        self.min_size = min_size;
247        self
248    }
249
250    /// Set the maximum number of entries.
251    pub fn max_size(mut self, max_size: usize) -> Self {
252        self.max_size = Some(max_size);
253        self
254    }
255}
256
257impl<K, V, KT, VT> Generator<HashMap<KT, VT>> for HashMapGenerator<K, V, KT, VT>
258where
259    K: Generator<KT>,
260    V: Generator<VT>,
261    KT: Eq + std::hash::Hash,
262{
263    fn do_draw(&self, tc: &TestCase) -> HashMap<KT, VT> {
264        if let Some(max) = self.max_size {
265            if self.min_size > max {
266                invalid_argument!("Cannot have max_size < min_size");
267            }
268        }
269        tc.start_span(labels::MAP);
270        let map = match self.enumerated_key_pool() {
271            Some(pool) => self.draw_from_key_pool(tc, pool),
272            None => self.draw_by_rejection(tc),
273        };
274        tc.stop_span(false);
275        map
276    }
277}
278
279impl<K, V, KT, VT> HashMapGenerator<K, V, KT, VT>
280where
281    K: Generator<KT>,
282    V: Generator<VT>,
283    KT: Eq + std::hash::Hash,
284{
285    /// The distinct values of an enumerable key generator, in first
286    /// occurrence order, when there are few enough of them to draw without
287    /// replacement. Mirrors [`HashSetGenerator::enumerated_pool`].
288    fn enumerated_key_pool(&self) -> Option<Vec<KT>> {
289        let values = self.keys.enumerate_values()?;
290        if values.is_empty() || values.len() > MAX_UNIQUE_POOL {
291            return None;
292        }
293        let mut by_hash: HashMap<u64, Vec<usize>> = HashMap::new();
294        let mut pool: Vec<KT> = Vec::new();
295        for v in values {
296            let bucket = by_hash.entry(fingerprint(&v)).or_default();
297            if bucket.iter().any(|&i| pool[i] == v) {
298                continue;
299            }
300            bucket.push(pool.len());
301            pool.push(v);
302        }
303        Some(pool)
304    }
305
306    /// Draw keys as indices into a shrinking pool of the remaining values,
307    /// avoiding the coupon-collector problem when the map must contain most
308    /// of a small key alphabet; each key's value is drawn right after it.
309    fn draw_from_key_pool(&self, tc: &TestCase, mut remaining: Vec<KT>) -> HashMap<KT, VT> {
310        if self.min_size > remaining.len() {
311            invalid_argument!(
312                "Cannot generate a map: min_size {} is larger than the {} distinct keys the key generator can produce",
313                self.min_size,
314                remaining.len()
315            );
316        }
317        let effective_max = self
318            .max_size
319            .map_or(remaining.len(), |m| m.min(remaining.len()));
320        let mut collection = Collection::new(tc, self.min_size, Some(effective_max));
321        let mut map = HashMap::new();
322        loop {
323            if remaining.is_empty() || !collection.more() {
324                break;
325            }
326            let j = tc.generate_integer_i64(0, remaining.len() as i64 - 1) as usize;
327            let key = remaining.remove(j);
328            let value = self.values.do_draw(tc);
329            map.insert(key, value);
330        }
331        map
332    }
333
334    fn draw_by_rejection(&self, tc: &TestCase) -> HashMap<KT, VT> {
335        let mut collection = Collection::new(tc, self.min_size, self.max_size);
336        let mut map = HashMap::new();
337        while collection.more() {
338            let key = self.keys.do_draw(tc);
339            match map.entry(key) {
340                std::collections::hash_map::Entry::Occupied(_) => {
341                    collection.reject(Some("duplicate key"));
342                }
343                std::collections::hash_map::Entry::Vacant(entry) => {
344                    let value = self.values.do_draw(tc);
345                    entry.insert(value);
346                }
347            }
348        }
349        hegel_internal_assert!(map.len() >= self.min_size);
350        map
351    }
352}
353
354/// Generate hash maps.
355///
356/// See [`HashMapGenerator`] for builder methods.
357///
358/// # Example
359///
360/// ```ignore
361/// use hegel::generators as gs;
362/// use std::collections::HashMap;
363///
364/// let map: HashMap<i32, String> = tc.draw(gs::hashmaps(gs::integers(), gs::text()));
365/// ```
366pub fn hashmaps<KT, VT, K: Generator<KT>, V: Generator<VT>>(
367    keys: K,
368    values: V,
369) -> HashMapGenerator<K, V, KT, VT> {
370    HashMapGenerator {
371        keys,
372        values,
373        min_size: 0,
374        max_size: None,
375        _phantom: PhantomData,
376    }
377}
378
379/// Generator for fixed-size arrays `[T; N]`. Created by [`arrays()`].
380pub struct ArrayGenerator<G, T, const N: usize> {
381    element: G,
382    _phantom: PhantomData<fn() -> T>,
383}
384
385impl<G, T, const N: usize> ArrayGenerator<G, T, N> {
386    #[doc(hidden)]
387    pub fn new(element: G) -> Self {
388        ArrayGenerator {
389            element,
390            _phantom: PhantomData,
391        }
392    }
393}
394
395/// Generate fixed-size arrays `[T; N]` with elements from the given generator.
396pub fn arrays<G: Generator<T> + Send + Sync, T, const N: usize>(
397    element: G,
398) -> ArrayGenerator<G, T, N> {
399    ArrayGenerator::new(element)
400}
401
402impl<G: Generator<T> + Send + Sync, T, const N: usize> Generator<[T; N]>
403    for ArrayGenerator<G, T, N>
404{
405    fn do_draw(&self, tc: &TestCase) -> [T; N] {
406        tc.start_span(labels::TUPLE);
407        let result = std::array::from_fn(|_| self.element.do_draw(tc));
408        tc.stop_span(false);
409        result
410    }
411}
hegel/generators/collections.rs

hegel/generators/
collections.rs
