crypto_primes/
generic.rs

1use rand_core::CryptoRngCore;
2
3#[cfg(feature = "multicore")]
4use rayon::iter::{ParallelBridge, ParallelIterator};
5
6use crate::SieveFactory;
7
8/// Sieves through the results of `sieve_factory` and returns the first item for which `predicate` is `true`.
9///
10/// If `sieve_factory` signals that no more results can be created, returns `None`.
11pub fn sieve_and_find<R, S>(
12    rng: &mut R,
13    sieve_factory: S,
14    predicate: impl Fn(&mut R, &S::Item) -> bool,
15) -> Option<S::Item>
16where
17    S: SieveFactory,
18    R: CryptoRngCore + ?Sized,
19{
20    // We could use `SieveIterator` here, but it requires cloning the `rng`.
21    // Unlike the parallel version, it is avoidable here.
22
23    let mut sieve_factory = sieve_factory;
24    let mut sieve = sieve_factory.make_sieve(rng, None)?;
25
26    loop {
27        if let Some(value) = sieve.find(|num| predicate(rng, num)) {
28            return Some(value);
29        }
30        if let Some(new_sieve) = sieve_factory.make_sieve(rng, Some(&sieve)) {
31            sieve = new_sieve;
32        } else {
33            return None;
34        }
35    }
36}
37
38/// Sieves through the results of `sieve_factory` using a thread pool with `threadcount` threads,
39/// and returns the first item for which `predicate` is `true`.
40///
41/// If `sieve_factory` signals that no more results can be created, returns `None`.
42#[cfg(feature = "multicore")]
43pub fn par_sieve_and_find<R, S, F>(rng: &mut R, sieve_factory: S, predicate: F, threadcount: usize) -> Option<S::Item>
44where
45    R: CryptoRngCore + Clone + Send + Sync,
46    S: Send + Sync + SieveFactory,
47    S::Sieve: Send,
48    S::Item: Send,
49    F: Sync + Fn(&mut R, &S::Item) -> bool,
50{
51    let threadpool = rayon::ThreadPoolBuilder::new()
52        .num_threads(threadcount)
53        .build()
54        .expect("If the platform can spawn threads, then this call will work.");
55
56    let mut iter_rng = rng.clone();
57    let iter = SieveIterator::new(&mut iter_rng, sieve_factory)?;
58
59    threadpool.install(|| {
60        iter.par_bridge().find_any(|c| {
61            let mut rng = rng.clone();
62            predicate(&mut rng, c)
63        })
64    })
65}
66
67/// A structure that chains the creation of sieves, returning the results from one until it is exhausted,
68/// and then creating a new one.
69#[derive(Debug)]
70pub struct SieveIterator<'a, R: CryptoRngCore + ?Sized, S: SieveFactory> {
71    sieve_factory: S,
72    sieve: S::Sieve,
73    rng: &'a mut R,
74}
75
76impl<'a, R: CryptoRngCore + ?Sized, S: SieveFactory> SieveIterator<'a, R, S> {
77    /// Creates a new chained iterator producing results from sieves returned from `sieve_factory`.
78    pub fn new(rng: &'a mut R, sieve_factory: S) -> Option<Self> {
79        let mut sieve_factory = sieve_factory;
80        let sieve = sieve_factory.make_sieve(rng, None)?;
81        Some(Self {
82            sieve_factory,
83            rng,
84            sieve,
85        })
86    }
87}
88
89impl<R: CryptoRngCore, S: SieveFactory> Iterator for SieveIterator<'_, R, S> {
90    type Item = S::Item;
91
92    fn next(&mut self) -> Option<Self::Item> {
93        loop {
94            if let Some(result) = self.sieve.next() {
95                return Some(result);
96            }
97
98            self.sieve = self.sieve_factory.make_sieve(self.rng, Some(&self.sieve))?;
99        }
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use rand_core::{CryptoRngCore, OsRng};
106
107    use super::sieve_and_find;
108    use crate::SieveFactory;
109
110    #[cfg(feature = "multicore")]
111    use super::par_sieve_and_find;
112
113    #[test]
114    fn test_exhaustable_sieve_factory() {
115        // Test the logic handling the case of the sieve factory not being able to produce new sieves.
116        struct TestSieveFactory {
117            count: usize,
118        }
119
120        impl SieveFactory for TestSieveFactory {
121            type Item = usize;
122            type Sieve = core::ops::Range<usize>;
123
124            fn make_sieve(
125                &mut self,
126                _rng: &mut (impl CryptoRngCore + ?Sized),
127                previous_sieve: Option<&Self::Sieve>,
128            ) -> Option<Self::Sieve> {
129                self.count += 1;
130                if previous_sieve.is_none() {
131                    Some(self.count * 10..(self.count * 10 + 2))
132                } else {
133                    None
134                }
135            }
136        }
137
138        let factory = TestSieveFactory { count: 0 };
139        let result = sieve_and_find(&mut OsRng, factory, |_rng, num| *num == 11);
140        assert!(result.is_some());
141
142        #[cfg(feature = "multicore")]
143        {
144            let factory = TestSieveFactory { count: 0 };
145            let result = par_sieve_and_find(&mut OsRng, factory, |_rng, num| *num == 11, 1);
146            assert!(result.is_some());
147        }
148
149        let factory = TestSieveFactory { count: 0 };
150        let result = sieve_and_find(&mut OsRng, factory, |_rng, num| *num == 20);
151        assert!(result.is_none());
152
153        #[cfg(feature = "multicore")]
154        {
155            let factory = TestSieveFactory { count: 0 };
156            let result = par_sieve_and_find(&mut OsRng, factory, |_rng, num| *num == 20, 1);
157            assert!(result.is_none());
158        }
159    }
160}
crypto_primes/generic.rs

crypto_primes/
generic.rs
