Generalize sieving

Author: fjarri

src/generic.rs

@@ -0,0 +1,95 @@
+use rand_core::CryptoRngCore;
+
+#[cfg(feature = "multicore")]
+use rayon::iter::{ParallelBridge, ParallelIterator};
+
+use crate::SieveFactory;
+
+/// Sieves through the results of `sieve_factory` and returns the first item for which `predicate` is `true`.
+///
+/// If `sieve_factory` signals that no more results can be created, returns `None`.
+pub fn sieve_and_find<R, S, T>(rng: &mut R, sieve_factory: &S, predicate: impl Fn(&mut R, &T) -> bool) -> Option<T>
+where
+    S: SieveFactory<T>,
+    R: CryptoRngCore,
+{
+    // We could use `SieveIterator` here, but it requires cloning the `rng`.
+    // Unlike the parallel version, it is avoidable here.
+
+    let mut sieve = sieve_factory.make_sieve(rng, None)?;
+
+    loop {
+        if let Some(value) = sieve.find(|num| predicate(rng, num)) {
+            return Some(value);
+        }
+        if let Some(new_sieve) = sieve_factory.make_sieve(rng, Some(&sieve)) {
+            sieve = new_sieve;
+        } else {
+            return None;
+        }
+    }
+}
+
+/// Sieves through the results of `sieve_factory` using a thread pool with `threadcount` threads,
+/// and returns the first item for which `predicate` is `true`.
+///
+/// If `sieve_factory` signals that no more results can be created, returns `None`.
+#[cfg(feature = "multicore")]
+pub fn par_sieve_and_find<R, S, T, F>(rng: &mut R, sieve_factory: S, predicate: F, threadcount: usize) -> Option<T>
+where
+    R: CryptoRngCore + Clone + Send + Sync,
+    T: Send,
+    S: Send + Sync + SieveFactory<T>,
+    S::Sieve: Send,
+    F: Sync + Fn(&mut R, &T) -> bool,
+{
+    let threadpool = rayon::ThreadPoolBuilder::new()
+        .num_threads(threadcount)
+        .build()
+        .expect("If the platform can spawn threads, then this call will work.");
+
+    let mut iter_rng = rng.clone();
+    let iter = SieveIterator::new(&mut iter_rng, sieve_factory)?;
+
+    threadpool.install(|| {
+        iter.par_bridge().find_any(|c| {
+            let mut rng = rng.clone();
+            predicate(&mut rng, c)
+        })
+    })
+}
+
+/// A structure that chains the creation of sieves, returning the results from one until it is exhausted,
+/// and then creating a new one.
+#[derive(Debug)]
+pub struct SieveIterator<'a, R: CryptoRngCore, T, S: SieveFactory<T>> {
+    sieve_factory: S,
+    sieve: S::Sieve,
+    rng: &'a mut R,
+}
+
+impl<'a, R: CryptoRngCore, T, S: SieveFactory<T>> SieveIterator<'a, R, T, S> {
+    /// Creates a new chained iterator producing results from sieves returned from `sieve_factory`.
+    pub fn new(rng: &'a mut R, sieve_factory: S) -> Option<Self> {
+        let sieve = sieve_factory.make_sieve(rng, None)?;
+        Some(Self {
+            sieve_factory,
+            rng,
+            sieve,
+        })
+    }
+}
+
+impl<'a, R: CryptoRngCore, T, S: SieveFactory<T>> Iterator for SieveIterator<'a, R, T, S> {
+    type Item = T;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            if let Some(result) = self.sieve.next() {
+                return Some(result);
+            }
+
+            self.sieve = self.sieve_factory.make_sieve(self.rng, Some(&self.sieve))?;
+        }
+    }
+}

src/hazmat.rs

@@ -14,7 +14,7 @@ mod sieve;
 
 pub use lucas::{lucas_test, AStarBase, BruteForceBase, LucasBase, LucasCheck, SelfridgeBase};
 pub use miller_rabin::MillerRabin;
-pub use sieve::{random_odd_integer, Sieve};
+pub use sieve::{random_odd_integer, DefaultSieveFactory, Sieve};
 
 /// Possible results of various primality tests.
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]

src/hazmat/sieve.rs

@@ -2,12 +2,13 @@
 //! before proceeding with slower tests.
 
 use alloc::{vec, vec::Vec};
-use core::num::NonZeroU32;
+use core::num::{NonZero, NonZeroU32};
 
 use crypto_bigint::{Integer, Odd, RandomBits};
 use rand_core::CryptoRngCore;
 
 use crate::hazmat::precomputed::{SmallPrime, LAST_SMALL_PRIME, RECIPROCALS, SMALL_PRIMES};
+use crate::traits::SieveFactory;
 
 /// Returns a random odd integer with given bit length
 /// (that is, with both `0` and `bit_length-1` bits set).
@@ -244,6 +245,42 @@ impl<T: Integer> Iterator for Sieve<T> {
     }
 }
 
+/// A sieve returning numbers that are not multiples of a set of small factors.
+#[derive(Debug, Clone, Copy)]
+pub struct DefaultSieveFactory {
+    max_bit_length: NonZeroU32,
+    safe_primes: bool,
+}
+
+impl DefaultSieveFactory {
+    /// Creates a factory that produces sieves returning numbers of `max_bit_length` bits (with the top bit set).
+    ///
+    /// If `safe_primes` is `true`, additionally filters out such `n` that `(n - 1) / 2` are divisible
+    /// by any of the small factors tested.
+    pub fn new(max_bit_length: u32, safe_primes: bool) -> Self {
+        let max_bit_length = NonZero::new(max_bit_length).expect("`bit_length` should be non-zero");
+        Self {
+            max_bit_length,
+            safe_primes,
+        }
+    }
+}
+
+impl<T: Integer + RandomBits> SieveFactory<T> for DefaultSieveFactory {
+    type Sieve = Sieve<T>;
+    fn make_sieve(&self, rng: &mut impl CryptoRngCore, _previous_sieve: Option<&Self::Sieve>) -> Option<Self::Sieve> {
+        if !self.safe_primes && self.max_bit_length.get() < 2 {
+            panic!("`bit_length` must be 2 or greater.");
+        }
+        if self.safe_primes && self.max_bit_length.get() < 3 {
+            panic!("`bit_length` must be 3 or greater.");
+        }
+
+        let start = random_odd_integer::<T>(rng, self.max_bit_length);
+        Some(Sieve::new(start.get(), self.max_bit_length, self.safe_primes))
+    }
+}
+
 #[cfg(test)]
 mod tests {
 

src/lib.rs

@@ -16,16 +16,21 @@
 
 extern crate alloc;
 
+mod generic;
 pub mod hazmat;
 mod presets;
 mod traits;
 
+pub use generic::{sieve_and_find, SieveIterator};
 pub use presets::{generate_prime_with_rng, generate_safe_prime_with_rng, is_prime_with_rng, is_safe_prime_with_rng};
-pub use traits::RandomPrimeWithRng;
+pub use traits::{RandomPrimeWithRng, SieveFactory};
 
 #[cfg(feature = "default-rng")]
 pub use presets::{generate_prime, generate_safe_prime, is_prime, is_safe_prime};
 #[cfg(all(feature = "default-rng", feature = "multicore"))]
 pub use presets::{par_generate_prime, par_generate_safe_prime};
 #[cfg(feature = "multicore")]
 pub use presets::{par_generate_prime_with_rng, par_generate_safe_prime_with_rng};
+
+#[cfg(feature = "multicore")]
+pub use generic::par_sieve_and_find;

src/presets.rs

@@ -1,15 +1,16 @@
-use core::num::NonZero;
-
 use crypto_bigint::{Integer, Limb, Odd, RandomBits, RandomMod};
 use rand_core::CryptoRngCore;
 
 #[cfg(feature = "default-rng")]
 use rand_core::OsRng;
 
-#[cfg(feature = "multicore")]
-use rayon::iter::{ParallelBridge, ParallelIterator};
+use crate::{
+    generic::sieve_and_find,
+    hazmat::{lucas_test, AStarBase, DefaultSieveFactory, LucasCheck, MillerRabin, Primality},
+};
 
-use crate::hazmat::{lucas_test, random_odd_integer, AStarBase, LucasCheck, MillerRabin, Primality, Sieve};
+#[cfg(feature = "multicore")]
+use crate::generic::par_sieve_and_find;
 
 /// Returns a random prime of size `bit_length` using [`OsRng`] as the RNG.
 ///
@@ -83,18 +84,8 @@ pub fn generate_prime_with_rng<T: Integer + RandomBits + RandomMod>(
     rng: &mut impl CryptoRngCore,
     bit_length: u32,
 ) -> T {
-    if bit_length < 2 {
-        panic!("`bit_length` must be 2 or greater.");
-    }
-    let bit_length = NonZero::new(bit_length).expect("`bit_length` should be non-zero");
-    // Empirically, this loop is traversed 1 time.
-    loop {
-        let start = random_odd_integer::<T>(rng, bit_length);
-        let mut sieve = Sieve::new(start.get(), bit_length, false);
-        if let Some(prime) = sieve.find(|num| is_prime_with_rng(rng, num)) {
-            return prime;
-        }
-    }
+    sieve_and_find(rng, &DefaultSieveFactory::new(bit_length, false), is_prime_with_rng)
+        .expect("will produce a result eventually")
 }
 
 /// Returns a random safe prime (that is, such that `(n - 1) / 2` is also prime)
@@ -107,19 +98,8 @@ pub fn generate_safe_prime_with_rng<T: Integer + RandomBits + RandomMod>(
     rng: &mut impl CryptoRngCore,
     bit_length: u32,
 ) -> T {
-    if bit_length < 3 {
-        panic!("`bit_length` must be 3 or greater.");
-    }
-    let bit_length = NonZero::new(bit_length).expect("`bit_length` should be non-zero");
-    loop {
-        let start = random_odd_integer::<T>(rng, bit_length);
-        let sieve = Sieve::new(start.get(), bit_length, true);
-        for num in sieve {
-            if is_safe_prime_with_rng(rng, &num) {
-                return num;
-            }
-        }
-    }
+    sieve_and_find(rng, &DefaultSieveFactory::new(bit_length, true), is_safe_prime_with_rng)
+        .expect("will produce a result eventually")
 }
 
 /// Returns a random prime of size `bit_length` using the provided RNG.
@@ -138,31 +118,13 @@ pub fn par_generate_prime_with_rng<T>(
 where
     T: Integer + RandomBits + RandomMod,
 {
-    if bit_length < 2 {
-        panic!("`bit_length` must be 2 or greater.");
-    }
-    let bit_length = NonZero::new(bit_length).expect("`bit_length` should be non-zero");
-
-    let threadpool = rayon::ThreadPoolBuilder::new()
-        .num_threads(threadcount)
-        .build()
-        .expect("If the platform can spawn threads, then this call will work.");
-    let start = random_odd_integer::<T>(rng, bit_length).get();
-    let sieve = Sieve::new(start, bit_length, false);
-
-    let prime = threadpool.install(|| {
-        sieve.par_bridge().find_any(|c| {
-            let mut rng = rng.clone();
-            is_prime_with_rng(&mut rng, c)
-        })
-    });
-    match prime {
-        Some(prime) => prime,
-        None => {
-            drop(threadpool);
-            par_generate_prime_with_rng(rng, bit_length.get(), threadcount)
-        }
-    }
+    par_sieve_and_find(
+        rng,
+        DefaultSieveFactory::new(bit_length, false),
+        is_prime_with_rng,
+        threadcount,
+    )
+    .expect("will produce a result eventually")
 }
 
 /// Returns a random safe prime (that is, such that `(n - 1) / 2` is also prime)
@@ -183,31 +145,13 @@ pub fn par_generate_safe_prime_with_rng<T>(
 where
     T: Integer + RandomBits + RandomMod,
 {
-    if bit_length < 2 {
-        panic!("`bit_length` must be 2 or greater.");
-    }
-    let bit_length = NonZero::new(bit_length).expect("`bit_length` should be non-zero");
-
-    let threadpool = rayon::ThreadPoolBuilder::new()
-        .num_threads(threadcount)
-        .build()
-        .expect("If the platform can spawn threads, then this call will work.");
-    let start = random_odd_integer::<T>(rng, bit_length).get();
-    let sieve = Sieve::new(start, bit_length, true);
-
-    let prime = threadpool.install(|| {
-        sieve.par_bridge().find_any(|c| {
-            let mut rng = rng.clone();
-            is_safe_prime_with_rng(&mut rng, c)
-        })
-    });
-    match prime {
-        Some(prime) => prime,
-        None => {
-            drop(threadpool);
-            par_generate_safe_prime_with_rng(rng, bit_length.get(), threadcount)
-        }
-    }
+    par_sieve_and_find(
+        rng,
+        DefaultSieveFactory::new(bit_length, true),
+        is_safe_prime_with_rng,
+        threadcount,
+    )
+    .expect("will produce a result eventually")
 }
 
 /// Probabilistically checks if the given number is prime using the provided RNG.

src/traits.rs

@@ -3,6 +3,17 @@ use rand_core::CryptoRngCore;
 
 use crate::{generate_prime_with_rng, generate_safe_prime_with_rng, is_prime_with_rng, is_safe_prime_with_rng};
 
+/// A type producing sieves for random prime generation.
+pub trait SieveFactory<T> {
+    /// The resulting sieve.
+    type Sieve: Iterator<Item = T>;
+
+    /// Makes a sieve given an RNG and the previous exhausted sieve (if any).
+    ///
+    /// Returning `None` signals that the prime generation should stop.
+    fn make_sieve(&self, rng: &mut impl CryptoRngCore, previous_sieve: Option<&Self::Sieve>) -> Option<Self::Sieve>;
+}
+
 /// Provides a generic way to access methods for random prime number generation
 /// and primality checking, wrapping the standalone functions ([`is_prime_with_rng`] etc).
 pub trait RandomPrimeWithRng {