protest_extras/generators/
collections.rs1use protest::{Generator, GeneratorConfig};
12use rand::Rng;
13use std::collections::HashMap;
14use std::hash::Hash;
15
16#[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 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 if value.len() > 1 {
65 shrinks.push(vec![value[0].clone()]);
66 }
67
68 if value.len() > self.min_len {
70 shrinks.push(value[..self.min_len].to_vec());
71 }
72
73 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#[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 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 if !value.is_empty() && self.min_len == 0 {
139 shrinks.push(vec![]);
140 }
141
142 if value.len() > self.min_len {
144 shrinks.push(value[..self.min_len].to_vec());
145 }
146
147 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 shrinks.push(shrunk);
156 }
157 }
158 }
159 }
160
161 Box::new(shrinks.into_iter())
162 }
163}
164
165#[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 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 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 if !value.is_empty() && self.min_len == 0 {
223 shrinks.push(vec![]);
224 }
225
226 if value.len() > self.min_len {
228 shrinks.push(value[..self.min_len].to_vec());
229 }
230
231 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#[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 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 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 if !value.is_empty() && self.min_size == 0 {
310 shrinks.push(HashMap::new());
311 }
312
313 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 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 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 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}