Skip to main content

protest_extras/generators/
numeric.rs

1//! Constrained numeric generators
2//!
3//! This module provides generators for numbers with specific constraints:
4//! - Positive integers
5//! - Even numbers
6//! - Prime numbers
7//! - Percentage values (0.0-100.0)
8//!
9//! All generators use std library only.
10
11use protest::{Generator, GeneratorConfig};
12use rand::Rng;
13
14// ============================================================================
15// Positive Integer Generator
16// ============================================================================
17
18/// Generator for positive integers (> 0)
19///
20/// Generic over any unsigned integer type
21#[derive(Debug, Clone)]
22pub struct PositiveIntGenerator<T> {
23    min: T,
24    max: T,
25}
26
27impl<T> PositiveIntGenerator<T>
28where
29    T: num_traits::PrimInt + num_traits::Unsigned + rand::distributions::uniform::SampleUniform,
30{
31    /// Create a new positive integer generator
32    pub fn new(min: T, max: T) -> Self {
33        Self { min, max }
34    }
35}
36
37impl<T> Generator<T> for PositiveIntGenerator<T>
38where
39    T: num_traits::PrimInt
40        + num_traits::Unsigned
41        + rand::distributions::uniform::SampleUniform
42        + Clone
43        + 'static,
44{
45    fn generate(&self, rng: &mut dyn rand::RngCore, _config: &GeneratorConfig) -> T {
46        rng.r#gen_range(self.min..=self.max)
47    }
48
49    fn shrink(&self, value: &T) -> Box<dyn Iterator<Item = T>> {
50        let mut shrinks = Vec::new();
51        let one = T::one();
52
53        // Try shrinking toward min
54        if *value > self.min {
55            shrinks.push(self.min);
56        }
57
58        // Try shrinking toward 1 (smallest positive)
59        if *value > one && one >= self.min {
60            shrinks.push(one);
61        }
62
63        // Try half
64        let two = one + one;
65        if *value > two {
66            shrinks.push(*value / two);
67        }
68
69        // Try value - 1
70        if *value > self.min {
71            shrinks.push(*value - one);
72        }
73
74        Box::new(shrinks.into_iter())
75    }
76}
77
78// ============================================================================
79// Even Number Generator
80// ============================================================================
81
82/// Generator for even numbers
83#[derive(Debug, Clone)]
84pub struct EvenNumberGenerator<T> {
85    min: T,
86    max: T,
87}
88
89impl<T> EvenNumberGenerator<T>
90where
91    T: num_traits::PrimInt,
92{
93    /// Create a new even number generator
94    ///
95    /// Note: min and max will be adjusted to the nearest even numbers if needed
96    pub fn new(min: T, max: T) -> Self {
97        let two = T::one() + T::one();
98
99        // Adjust min to nearest even number >= min
100        let adjusted_min = if min % two == T::zero() {
101            min
102        } else {
103            min + T::one()
104        };
105
106        // Adjust max to nearest even number <= max
107        let adjusted_max = if max % two == T::zero() {
108            max
109        } else {
110            max - T::one()
111        };
112
113        Self {
114            min: adjusted_min,
115            max: adjusted_max,
116        }
117    }
118}
119
120impl<T> Generator<T> for EvenNumberGenerator<T>
121where
122    T: num_traits::PrimInt + rand::distributions::uniform::SampleUniform + Clone + 'static,
123{
124    fn generate(&self, rng: &mut dyn rand::RngCore, _config: &GeneratorConfig) -> T {
125        let two = T::one() + T::one();
126        let range = (self.max - self.min) / two + T::one();
127        let offset = rng.r#gen_range(T::zero()..range);
128        self.min + (offset * two)
129    }
130
131    fn shrink(&self, value: &T) -> Box<dyn Iterator<Item = T>> {
132        let mut shrinks = Vec::new();
133        let two = T::one() + T::one();
134
135        // Try zero if in range
136        if T::zero() >= self.min && T::zero() <= self.max && *value != T::zero() {
137            shrinks.push(T::zero());
138        }
139
140        // Try min
141        if *value > self.min {
142            shrinks.push(self.min);
143        }
144
145        // Try half (rounded to even)
146        if *value > two {
147            let half = *value / two;
148            let half_even = if half % two == T::zero() {
149                half
150            } else {
151                half - T::one()
152            };
153            if half_even >= self.min {
154                shrinks.push(half_even);
155            }
156        }
157
158        // Try value - 2
159        if *value > self.min + two {
160            shrinks.push(*value - two);
161        }
162
163        Box::new(shrinks.into_iter())
164    }
165}
166
167// ============================================================================
168// Prime Number Generator
169// ============================================================================
170
171/// Generator for prime numbers
172#[derive(Debug, Clone)]
173pub struct PrimeNumberGenerator {
174    min: u64,
175    #[allow(dead_code)]
176    max: u64,
177    primes_cache: Vec<u64>,
178}
179
180impl PrimeNumberGenerator {
181    /// Create a new prime number generator
182    pub fn new(min: u64, max: u64) -> Self {
183        let primes = Self::generate_primes_in_range(min, max);
184        Self {
185            min,
186            max,
187            primes_cache: primes,
188        }
189    }
190
191    fn is_prime(n: u64) -> bool {
192        if n < 2 {
193            return false;
194        }
195        if n == 2 {
196            return true;
197        }
198        if n.is_multiple_of(2) {
199            return false;
200        }
201
202        let sqrt = (n as f64).sqrt() as u64;
203        for i in (3..=sqrt).step_by(2) {
204            if n.is_multiple_of(i) {
205                return false;
206            }
207        }
208        true
209    }
210
211    fn generate_primes_in_range(min: u64, max: u64) -> Vec<u64> {
212        let start = if min <= 2 { 2 } else { min };
213        (start..=max).filter(|&n| Self::is_prime(n)).collect()
214    }
215}
216
217impl Generator<u64> for PrimeNumberGenerator {
218    fn generate(&self, rng: &mut dyn rand::RngCore, _config: &GeneratorConfig) -> u64 {
219        if self.primes_cache.is_empty() {
220            // No primes in range, return min
221            self.min
222        } else {
223            let idx = rng.r#gen_range(0..self.primes_cache.len());
224            self.primes_cache[idx]
225        }
226    }
227
228    fn shrink(&self, value: &u64) -> Box<dyn Iterator<Item = u64>> {
229        let mut shrinks = Vec::new();
230
231        // Try smaller primes
232        for &prime in &self.primes_cache {
233            if prime < *value {
234                shrinks.push(prime);
235                if shrinks.len() >= 5 {
236                    break;
237                }
238            }
239        }
240
241        Box::new(shrinks.into_iter())
242    }
243}
244
245// ============================================================================
246// Percentage Generator
247// ============================================================================
248
249/// Generator for percentage values (0.0 to 100.0)
250#[derive(Debug, Clone, Copy)]
251pub struct PercentageGenerator {
252    allow_zero: bool,
253    allow_hundred: bool,
254}
255
256impl PercentageGenerator {
257    /// Create a new percentage generator (0.0 to 100.0 inclusive)
258    pub fn new() -> Self {
259        Self {
260            allow_zero: true,
261            allow_hundred: true,
262        }
263    }
264
265    /// Create a percentage generator excluding 0.0
266    pub fn positive() -> Self {
267        Self {
268            allow_zero: false,
269            allow_hundred: true,
270        }
271    }
272
273    /// Create a percentage generator with custom bounds
274    pub fn with_bounds(allow_zero: bool, allow_hundred: bool) -> Self {
275        Self {
276            allow_zero,
277            allow_hundred,
278        }
279    }
280}
281
282impl Default for PercentageGenerator {
283    fn default() -> Self {
284        Self::new()
285    }
286}
287
288impl Generator<f64> for PercentageGenerator {
289    fn generate(&self, rng: &mut dyn rand::RngCore, _config: &GeneratorConfig) -> f64 {
290        let min = if self.allow_zero { 0.0 } else { 0.01 };
291        let max = if self.allow_hundred { 100.0 } else { 99.99 };
292
293        rng.r#gen_range(min..=max)
294    }
295
296    fn shrink(&self, value: &f64) -> Box<dyn Iterator<Item = f64>> {
297        let mut shrinks = Vec::new();
298
299        // Try common values
300        if self.allow_zero && *value > 0.0 {
301            shrinks.push(0.0);
302        }
303
304        if *value > 1.0 {
305            shrinks.push(1.0);
306        }
307
308        if *value > 50.0 {
309            shrinks.push(50.0);
310        }
311
312        if self.allow_hundred && *value < 100.0 {
313            shrinks.push(100.0);
314        }
315
316        // Try half
317        if *value > 1.0 {
318            shrinks.push(*value / 2.0);
319        }
320
321        Box::new(shrinks.into_iter())
322    }
323}
324
325#[cfg(test)]
326mod tests {
327    use super::*;
328    use rand::thread_rng;
329
330    #[test]
331    fn test_positive_int_generator() {
332        let generator = PositiveIntGenerator::new(1u32, 100u32);
333        let mut rng = thread_rng();
334        let config = GeneratorConfig::default();
335
336        for _ in 0..20 {
337            let n = generator.generate(&mut rng, &config);
338            assert!((1..=100).contains(&n));
339            assert!(n > 0);
340        }
341    }
342
343    #[test]
344    fn test_even_number_generator() {
345        let generator = EvenNumberGenerator::new(0i32, 100i32);
346        let mut rng = thread_rng();
347        let config = GeneratorConfig::default();
348
349        for _ in 0..20 {
350            let n = generator.generate(&mut rng, &config);
351            assert!((0..=100).contains(&n));
352            assert_eq!(n % 2, 0, "Generated number {} should be even", n);
353        }
354    }
355
356    #[test]
357    fn test_even_number_generator_odd_bounds() {
358        // Test that odd bounds get adjusted correctly
359        let generator = EvenNumberGenerator::new(1i32, 99i32);
360        let mut rng = thread_rng();
361        let config = GeneratorConfig::default();
362
363        for _ in 0..20 {
364            let n = generator.generate(&mut rng, &config);
365            assert!((2..=98).contains(&n)); // Adjusted to 2..98
366            assert_eq!(n % 2, 0);
367        }
368    }
369
370    #[test]
371    fn test_prime_number_generator() {
372        let generator = PrimeNumberGenerator::new(2, 50);
373        let mut rng = thread_rng();
374        let config = GeneratorConfig::default();
375
376        for _ in 0..20 {
377            let n = generator.generate(&mut rng, &config);
378            assert!((2..=50).contains(&n));
379            assert!(PrimeNumberGenerator::is_prime(n));
380        }
381    }
382
383    #[test]
384    fn test_prime_number_is_prime() {
385        assert!(PrimeNumberGenerator::is_prime(2));
386        assert!(PrimeNumberGenerator::is_prime(3));
387        assert!(PrimeNumberGenerator::is_prime(5));
388        assert!(PrimeNumberGenerator::is_prime(7));
389        assert!(PrimeNumberGenerator::is_prime(11));
390        assert!(PrimeNumberGenerator::is_prime(13));
391        assert!(PrimeNumberGenerator::is_prime(97));
392
393        assert!(!PrimeNumberGenerator::is_prime(0));
394        assert!(!PrimeNumberGenerator::is_prime(1));
395        assert!(!PrimeNumberGenerator::is_prime(4));
396        assert!(!PrimeNumberGenerator::is_prime(6));
397        assert!(!PrimeNumberGenerator::is_prime(8));
398        assert!(!PrimeNumberGenerator::is_prime(9));
399        assert!(!PrimeNumberGenerator::is_prime(100));
400    }
401
402    #[test]
403    fn test_percentage_generator() {
404        let generator = PercentageGenerator::new();
405        let mut rng = thread_rng();
406        let config = GeneratorConfig::default();
407
408        for _ in 0..20 {
409            let p = generator.generate(&mut rng, &config);
410            assert!((0.0..=100.0).contains(&p));
411        }
412    }
413
414    #[test]
415    fn test_percentage_generator_positive() {
416        let generator = PercentageGenerator::positive();
417        let mut rng = thread_rng();
418        let config = GeneratorConfig::default();
419
420        for _ in 0..20 {
421            let p = generator.generate(&mut rng, &config);
422            assert!(p > 0.0 && p <= 100.0);
423        }
424    }
425}