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
12pub 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
44pub 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 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}