Skip to main content

protest_extras/generators/
collections.rs

1//! Specialized collection generators
2//!
3//! This module provides generators for collections with specific properties:
4//! - Non-empty vectors (guaranteed length >= 1)
5//! - Sorted collections
6//! - Unique element collections
7//! - Bounded size maps
8//!
9//! All generators use std library only.
10
11use protest::{Generator, GeneratorConfig};
12use rand::Rng;
13use std::collections::HashMap;
14use std::hash::Hash;
15
16// ============================================================================
17// NonEmpty Vec Generator
18// ============================================================================
19
20/// Generator for non-empty vectors
21///
22/// Guarantees that generated vectors always have at least one element
23#[derive(Debug, Clone)]
24pub struct NonEmptyVecGenerator<T, G> {
25    element_generator: G,
26    min_len: usize,
27    max_len: usize,
28    _phantom: std::marker::PhantomData<T>,
29}
30
31impl<T, G> NonEmptyVecGenerator<T, G>
32where
33    G: Generator<T>,
34{
35    /// Create a new non-empty vec generator
36    ///
37    /// min_len will be adjusted to at least 1
38    pub fn new(element_generator: G, min_len: usize, max_len: usize) -> Self {
39        Self {
40            element_generator,
41            min_len: min_len.max(1),
42            max_len: max_len.max(1),
43            _phantom: std::marker::PhantomData,
44        }
45    }
46}
47
48impl<T, G> Generator<Vec<T>> for NonEmptyVecGenerator<T, G>
49where
50    T: Clone + 'static,
51    G: Generator<T> + Clone + 'static,
52{
53    fn generate(&self, rng: &mut dyn rand::RngCore, config: &GeneratorConfig) -> Vec<T> {
54        let len = rng.r#gen_range(self.min_len..=self.max_len);
55        (0..len)
56            .map(|_| self.element_generator.generate(rng, config))
57            .collect()
58    }
59
60    fn shrink(&self, value: &Vec<T>) -> Box<dyn Iterator<Item = Vec<T>>> {
61        let mut shrinks = Vec::new();
62
63        // Try minimal size (1 element)
64        if value.len() > 1 {
65            shrinks.push(vec![value[0].clone()]);
66        }
67
68        // Try removing elements from end
69        if value.len() > self.min_len {
70            shrinks.push(value[..self.min_len].to_vec());
71        }
72
73        // Try removing one element at a time
74        if value.len() > self.min_len {
75            for i in 0..value.len().min(3) {
76                if i < value.len() {
77                    let mut shrunk = value.clone();
78                    shrunk.remove(i);
79                    if shrunk.len() >= self.min_len {
80                        shrinks.push(shrunk);
81                    }
82                }
83            }
84        }
85
86        Box::new(shrinks.into_iter())
87    }
88}
89
90// ============================================================================
91// Sorted Vec Generator
92// ============================================================================
93
94/// Generator for sorted vectors
95///
96/// Generates vectors that are already sorted in ascending order
97#[derive(Debug, Clone)]
98pub struct SortedVecGenerator<T, G> {
99    element_generator: G,
100    min_len: usize,
101    max_len: usize,
102    _phantom: std::marker::PhantomData<T>,
103}
104
105impl<T, G> SortedVecGenerator<T, G>
106where
107    G: Generator<T>,
108{
109    /// Create a new sorted vec generator
110    pub fn new(element_generator: G, min_len: usize, max_len: usize) -> Self {
111        Self {
112            element_generator,
113            min_len,
114            max_len,
115            _phantom: std::marker::PhantomData,
116        }
117    }
118}
119
120impl<T, G> Generator<Vec<T>> for SortedVecGenerator<T, G>
121where
122    T: Clone + Ord + 'static,
123    G: Generator<T> + Clone + 'static,
124{
125    fn generate(&self, rng: &mut dyn rand::RngCore, config: &GeneratorConfig) -> Vec<T> {
126        let len = rng.r#gen_range(self.min_len..=self.max_len);
127        let mut vec: Vec<T> = (0..len)
128            .map(|_| self.element_generator.generate(rng, config))
129            .collect();
130        vec.sort();
131        vec
132    }
133
134    fn shrink(&self, value: &Vec<T>) -> Box<dyn Iterator<Item = Vec<T>>> {
135        let mut shrinks = Vec::new();
136
137        // Try empty if allowed
138        if !value.is_empty() && self.min_len == 0 {
139            shrinks.push(vec![]);
140        }
141
142        // Try shrinking to min length
143        if value.len() > self.min_len {
144            shrinks.push(value[..self.min_len].to_vec());
145        }
146
147        // Try removing elements
148        if value.len() > self.min_len {
149            for i in 0..value.len().min(3) {
150                if i < value.len() {
151                    let mut shrunk = value.clone();
152                    shrunk.remove(i);
153                    if shrunk.len() >= self.min_len {
154                        // Should still be sorted after removal
155                        shrinks.push(shrunk);
156                    }
157                }
158            }
159        }
160
161        Box::new(shrinks.into_iter())
162    }
163}
164
165// ============================================================================
166// Unique Vec Generator
167// ============================================================================
168
169/// Generator for vectors with unique elements (no duplicates)
170#[derive(Debug, Clone)]
171pub struct UniqueVecGenerator<T, G> {
172    element_generator: G,
173    min_len: usize,
174    max_len: usize,
175    _phantom: std::marker::PhantomData<T>,
176}
177
178impl<T, G> UniqueVecGenerator<T, G>
179where
180    G: Generator<T>,
181{
182    /// Create a new unique vec generator
183    pub fn new(element_generator: G, min_len: usize, max_len: usize) -> Self {
184        Self {
185            element_generator,
186            min_len,
187            max_len,
188            _phantom: std::marker::PhantomData,
189        }
190    }
191}
192
193impl<T, G> Generator<Vec<T>> for UniqueVecGenerator<T, G>
194where
195    T: Clone + Eq + Hash + 'static,
196    G: Generator<T> + Clone + 'static,
197{
198    fn generate(&self, rng: &mut dyn rand::RngCore, config: &GeneratorConfig) -> Vec<T> {
199        let target_len = rng.r#gen_range(self.min_len..=self.max_len);
200        let mut vec = Vec::new();
201        let mut seen = std::collections::HashSet::new();
202
203        // Try to generate unique elements
204        let mut attempts = 0;
205        let max_attempts = target_len * 10;
206
207        while vec.len() < target_len && attempts < max_attempts {
208            let elem = self.element_generator.generate(rng, config);
209            if seen.insert(elem.clone()) {
210                vec.push(elem);
211            }
212            attempts += 1;
213        }
214
215        vec
216    }
217
218    fn shrink(&self, value: &Vec<T>) -> Box<dyn Iterator<Item = Vec<T>>> {
219        let mut shrinks = Vec::new();
220
221        // Try empty if allowed
222        if !value.is_empty() && self.min_len == 0 {
223            shrinks.push(vec![]);
224        }
225
226        // Try shrinking to min length
227        if value.len() > self.min_len {
228            shrinks.push(value[..self.min_len].to_vec());
229        }
230
231        // Try removing one element at a time
232        if value.len() > self.min_len {
233            for i in 0..value.len().min(3) {
234                if i < value.len() {
235                    let mut shrunk = value.clone();
236                    shrunk.remove(i);
237                    if shrunk.len() >= self.min_len {
238                        shrinks.push(shrunk);
239                    }
240                }
241            }
242        }
243
244        Box::new(shrinks.into_iter())
245    }
246}
247
248// ============================================================================
249// Bounded Map Generator
250// ============================================================================
251
252/// Generator for HashMap with size bounds
253#[derive(Debug, Clone)]
254pub struct BoundedMapGenerator<K, V, KG, VG> {
255    key_generator: KG,
256    value_generator: VG,
257    min_size: usize,
258    max_size: usize,
259    _phantom: std::marker::PhantomData<(K, V)>,
260}
261
262impl<K, V, KG, VG> BoundedMapGenerator<K, V, KG, VG>
263where
264    KG: Generator<K>,
265    VG: Generator<V>,
266{
267    /// Create a new bounded map generator
268    pub fn new(key_generator: KG, value_generator: VG, min_size: usize, max_size: usize) -> Self {
269        Self {
270            key_generator,
271            value_generator,
272            min_size,
273            max_size,
274            _phantom: std::marker::PhantomData,
275        }
276    }
277}
278
279impl<K, V, KG, VG> Generator<HashMap<K, V>> for BoundedMapGenerator<K, V, KG, VG>
280where
281    K: Clone + Eq + Hash + 'static,
282    V: Clone + 'static,
283    KG: Generator<K> + Clone + 'static,
284    VG: Generator<V> + Clone + 'static,
285{
286    fn generate(&self, rng: &mut dyn rand::RngCore, config: &GeneratorConfig) -> HashMap<K, V> {
287        use rand::Rng;
288        let target_size = rng.r#gen_range(self.min_size..=self.max_size);
289        let mut map = HashMap::new();
290
291        // Try to generate unique keys
292        let mut attempts = 0;
293        let max_attempts = target_size * 10;
294
295        while map.len() < target_size && attempts < max_attempts {
296            let key = self.key_generator.generate(rng, config);
297            let value = self.value_generator.generate(rng, config);
298            map.insert(key, value);
299            attempts += 1;
300        }
301
302        map
303    }
304
305    fn shrink(&self, value: &HashMap<K, V>) -> Box<dyn Iterator<Item = HashMap<K, V>>> {
306        let mut shrinks = Vec::new();
307
308        // Try empty if allowed
309        if !value.is_empty() && self.min_size == 0 {
310            shrinks.push(HashMap::new());
311        }
312
313        // Try removing entries
314        if value.len() > self.min_size {
315            let keys: Vec<_> = value.keys().cloned().collect();
316            for key in keys.iter().take(3) {
317                let mut shrunk = value.clone();
318                shrunk.remove(key);
319                if shrunk.len() >= self.min_size {
320                    shrinks.push(shrunk);
321                }
322            }
323        }
324
325        Box::new(shrinks.into_iter())
326    }
327}
328
329#[cfg(test)]
330mod tests {
331    use super::*;
332    use protest::IntGenerator;
333    use rand::thread_rng;
334
335    #[test]
336    fn test_non_empty_vec_generator() {
337        let generator = NonEmptyVecGenerator::new(IntGenerator::new(1, 100), 1, 10);
338        let mut rng = thread_rng();
339        let config = GeneratorConfig::default();
340
341        for _ in 0..10 {
342            let vec = generator.generate(&mut rng, &config);
343            assert!(!vec.is_empty(), "Vec should not be empty");
344            assert!(!vec.is_empty() && vec.len() <= 10);
345            for &elem in &vec {
346                assert!((1..=100).contains(&elem));
347            }
348        }
349    }
350
351    #[test]
352    fn test_non_empty_vec_min_size() {
353        // Even with min_len = 0, should have at least 1 element
354        let generator = NonEmptyVecGenerator::new(IntGenerator::new(1, 100), 0, 5);
355        let mut rng = thread_rng();
356        let config = GeneratorConfig::default();
357
358        for _ in 0..10 {
359            let vec = generator.generate(&mut rng, &config);
360            assert!(!vec.is_empty());
361        }
362    }
363
364    #[test]
365    fn test_sorted_vec_generator() {
366        let generator = SortedVecGenerator::new(IntGenerator::new(1, 100), 3, 10);
367        let mut rng = thread_rng();
368        let config = GeneratorConfig::default();
369
370        for _ in 0..10 {
371            let vec = generator.generate(&mut rng, &config);
372            assert!(vec.len() >= 3 && vec.len() <= 10);
373
374            // Check if sorted
375            for window in vec.windows(2) {
376                assert!(window[0] <= window[1], "Vec should be sorted: {:?}", vec);
377            }
378        }
379    }
380
381    #[test]
382    fn test_unique_vec_generator() {
383        let generator = UniqueVecGenerator::new(IntGenerator::new(1, 50), 3, 10);
384        let mut rng = thread_rng();
385        let config = GeneratorConfig::default();
386
387        for _ in 0..10 {
388            let vec = generator.generate(&mut rng, &config);
389            assert!(vec.len() >= 3 && vec.len() <= 10);
390
391            // Check for uniqueness
392            let mut seen = std::collections::HashSet::new();
393            for elem in &vec {
394                assert!(seen.insert(*elem), "Duplicate found: {} in {:?}", elem, vec);
395            }
396        }
397    }
398
399    #[test]
400    fn test_bounded_map_generator() {
401        let generator =
402            BoundedMapGenerator::new(IntGenerator::new(1, 100), IntGenerator::new(1, 100), 2, 8);
403        let mut rng = thread_rng();
404        let config = GeneratorConfig::default();
405
406        for _ in 0..10 {
407            let map = generator.generate(&mut rng, &config);
408            assert!(map.len() >= 2 && map.len() <= 8);
409
410            for (key, value) in &map {
411                assert!(*key >= 1 && *key <= 100);
412                assert!(*value >= 1 && *value <= 100);
413            }
414        }
415    }
416}