light_utils/
rand.rs

1use std::ops::{Bound, RangeBounds};
2
3use rand::{
4    distributions::uniform::{SampleRange, SampleUniform},
5    Rng,
6};
7
8use crate::prime::find_next_prime;
9
10const PRIME_RETRIES: usize = 10;
11
12/// Generates a random prime number in the given range. It returns `None` when
13/// generating such number was not possible.
14pub fn gen_prime<N, R, T>(rng: &mut N, range: R) -> Option<T>
15where
16    N: Rng,
17    R: Clone + RangeBounds<T> + SampleRange<T>,
18    T: Into<u32> + From<u32> + Copy + PartialOrd + SampleUniform,
19{
20    for _ in 0..PRIME_RETRIES {
21        let sample: T = rng.gen_range(range.clone());
22        let next_prime = find_next_prime(sample.into());
23
24        match range.end_bound() {
25            Bound::Included(end) => {
26                if next_prime > (*end).into() {
27                    continue;
28                }
29            }
30            Bound::Excluded(end) => {
31                if next_prime >= (*end).into() {
32                    continue;
33                }
34            }
35            _ => {}
36        };
37
38        return Some(T::from(next_prime));
39    }
40
41    None
42}
43
44/// Generates a random value in the given range, excluding the values provided
45/// in `exclude`.
46pub fn gen_range_exclude<N, R, T>(rng: &mut N, range: R, exclude: &[T]) -> T
47where
48    N: Rng,
49    R: Clone + SampleRange<T>,
50    T: PartialEq + SampleUniform,
51{
52    loop {
53        // This utility is supposed to be used only in unit tests. This `clone`
54        // is harmless and necessary (can't pass a reference to range, it has
55        // to be moved).
56        let sample = rng.gen_range(range.clone());
57        if !exclude.contains(&sample) {
58            return sample;
59        }
60    }
61}
62
63#[cfg(test)]
64mod test {
65    use rand::Rng;
66
67    use crate::prime::is_prime;
68
69    use super::*;
70
71    #[test]
72    fn test_gen_prime() {
73        let mut rng = rand::thread_rng();
74
75        let mut successful_gens = 0;
76        for i in 0..10_000 {
77            let sample: Option<u32> = gen_prime(&mut rng, 1..10_000);
78            println!("sample {i}: {sample:?}");
79            if let Some(sample) = sample {
80                successful_gens += 1;
81                assert!(is_prime(sample));
82            }
83        }
84
85        println!("generated {successful_gens} prime numbers out of 10000 iterations");
86    }
87
88    #[test]
89    fn test_gen_range_exclude() {
90        let mut rng = rand::thread_rng();
91
92        for n_excluded in 1..100 {
93            let excluded: Vec<u64> = (0..n_excluded).map(|_| rng.gen_range(0..100)).collect();
94
95            for _ in 0..10_000 {
96                let sample = gen_range_exclude(&mut rng, 0..100, excluded.as_slice());
97                for excluded in excluded.iter() {
98                    assert_ne!(&sample, excluded);
99                }
100            }
101        }
102    }
103}