Generalize sieving

src/generic.rs

@@ -0,0 +1,26 @@
+use rand_core::CryptoRngCore;
+
+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,
+{
+    let mut sieve = sieve_factory.make_sieve(rng, None)?;
+
+    loop {
+        let prime = sieve.find(|num| predicate(rng, num));
+        if prime.is_some() {
+            return prime;
+        }
+        if let Some(new_sieve) = sieve_factory.make_sieve(rng, Some(sieve)) {
+            sieve = new_sieve;
+        } else {
+            return None;
+        }
+    }
+}

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,12 +16,14 @@
 
 extern crate alloc;
 
+mod generic;
 pub mod hazmat;
 mod presets;
 mod traits;
 
+pub use generic::sieve_and_find;
 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};

src/presets.rs

@@ -1,3 +1,4 @@
+#[cfg(feature = "multicore")]
 use core::num::NonZero;
 
 use crypto_bigint::{Integer, Limb, Odd, RandomBits, RandomMod};
@@ -9,7 +10,13 @@ use rand_core::OsRng;
 #[cfg(feature = "multicore")]
 use rayon::iter::{ParallelBridge, ParallelIterator};
 
-use crate::hazmat::{lucas_test, random_odd_integer, AStarBase, LucasCheck, MillerRabin, Primality, Sieve};
+use crate::{
+    generic::sieve_and_find,
+    hazmat::{lucas_test, AStarBase, DefaultSieveFactory, LucasCheck, MillerRabin, Primality},
+};
+
+#[cfg(feature = "multicore")]
+use crate::hazmat::{random_odd_integer, Sieve};
 
 /// Returns a random prime of size `bit_length` using [`OsRng`] as the RNG.
 ///
@@ -83,18 +90,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 +104,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.

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 shoulds 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 {