Fix a bug where primes_lt and sieve_lt could panic if the underlying segmented sieve happend to run down past 0. (#55)

* Add tests of primes_geq, sieve_lt and primes_lt

* Make sieve_segment return an error when upper_limit is less than N, and use that signal in sieve_lt and primes_lt

* Add test of sieve_segment error

Author: JSorngard

src/generate.rs

@@ -194,12 +194,15 @@ pub const fn primes_lt<const N: usize, const MEM: usize>(
         // This is the smallest prime we have found so far.
         let mut smallest_found_prime = primes[N - 1 - total_primes_found];
         // Sieve for primes in the segment.
-        let upper_sieve: [bool; MEM] = sieve_segment(&base_sieve, upper_limit);
+        let (offset, upper_sieve) = match sieve_segment(&base_sieve, upper_limit) {
+            Ok(res) => (0, res),
+            Err(_) => ((MEM as u64 - upper_limit) as usize, base_sieve),
+        };
 
         let mut i: usize = 0;
-        while i < MEM {
+        while i < MEM - offset {
             // Iterate backwards through the upper sieve.
-            if upper_sieve[MEM - 1 - i] {
+            if upper_sieve[MEM - 1 - i - offset] {
                 smallest_found_prime = upper_limit - 1 - i as u64;
                 // Write every found prime to the primes array.
                 primes[N - 1 - total_primes_found] = smallest_found_prime;
@@ -368,7 +371,10 @@ pub const fn primes_geq<const N: usize, const MEM: usize>(
     let base_sieve: [bool; MEM] = sieve();
     let mut sieve_limit = lower_limit;
     'generate: while total_found_primes < N {
-        let upper_sieve = sieve_segment(&base_sieve, sieve_limit + mem64);
+        let upper_sieve = match sieve_segment(&base_sieve, sieve_limit + mem64) {
+            Ok(res) => res,
+            Err(_) => panic!("can not happen since we set upper limit to mem + nonzero stuff"),
+        };
 
         let mut i = 0;
         while i < MEM {
@@ -463,6 +469,36 @@ mod test {
         for &prime in primes_geq::<2_000, 2_008>(3_998_000).unwrap().as_slice() {
             assert!(is_prime(prime));
         }
+        assert_eq!(primes_geq::<0, 0>(10), Ok([]));
+        assert_eq!(primes_geq::<3, 3>(2), Ok([2, 3, 5]));
+        assert_eq!(
+            primes_geq::<3, 3>(10),
+            Err(GenerationError::TooSmallSieveSize)
+        );
+        assert_eq!(primes_geq::<2, 2>(3), Err(GenerationError::SieveOverrun(4)));
+    }
+
+    #[test]
+    fn sanity_check_primes_lt() {
+        {
+            const P: Result<[u64; 5], GenerationError> = primes_lt::<5, 5>(20);
+            assert_eq!(P, Ok([7, 11, 13, 17, 19]));
+        }
+        {
+            const P: Result<[u64; 5], GenerationError> = primes_lt::<5, 5>(12);
+            assert_eq!(P, Ok([2, 3, 5, 7, 11]));
+        }
+        {
+            const P: Result<[u64; 1], GenerationError> = primes_lt::<1, 2>(3);
+            assert_eq!(P, Ok([2]));
+        }
+        assert_eq!(primes_lt::<2, 2>(2), Err(GenerationError::TooSmallLimit));
+        assert_eq!(
+            primes_lt::<2, 2>(5),
+            Err(GenerationError::TooSmallSieveSize)
+        );
+        assert_eq!(primes_lt::<0, 2>(3), Ok([]));
+        assert_eq!(primes_lt::<3, 5>(4), Err(GenerationError::OutOfPrimes));
     }
 
     #[test]

src/sieve.rs

@@ -4,21 +4,40 @@ use core::fmt;
 
 use crate::isqrt;
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) struct SegmentedSieveError;
+
+impl fmt::Display for SegmentedSieveError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "the upper limit was smaller than `N`")
+    }
+}
+
+#[cfg(feature = "std")]
+impl std::error::Error for SegmentedSieveError {}
+
 /// Uses the primalities of the first `N` integers in `base_sieve` to sieve the numbers in the range `[upper_limit - N, upper_limit)`.
 /// Assumes that the base sieve contains the prime status of the `N` fist integers. The output is only meaningful
-/// for the numbers below `N^2`. Fails to compile if `N` is 0.
+/// for the numbers below `N^2`.
+///
+/// # Errors
+///
+/// Returns an error if `upper_limit` < `N`
 #[must_use = "the function only returns a new value and does not modify its inputs"]
 pub(crate) const fn sieve_segment<const N: usize>(
     base_sieve: &[bool; N],
     upper_limit: u64,
-) -> [bool; N] {
+) -> Result<[bool; N], SegmentedSieveError> {
     let mut segment_sieve = [true; N];
 
-    let lower_limit = upper_limit - N as u64;
+    let lower_limit = match upper_limit.checked_sub(N as u64) {
+        Some(diff) => diff,
+        None => return Err(SegmentedSieveError),
+    };
 
     if lower_limit == 0 && N > 1 {
         // If the lower limit is 0 we can just return the base sieve.
-        return *base_sieve;
+        return Ok(*base_sieve);
     } else if lower_limit == 1 && N > 0 {
         // In case 1 is included in the upper sieve we need to treat it as a special case
         // since it's not a multiple of any prime in `base_sieve` even though it's not prime.
@@ -48,7 +67,7 @@ pub(crate) const fn sieve_segment<const N: usize>(
         i += 1;
     }
 
-    segment_sieve
+    Ok(segment_sieve)
 }
 
 /// Returns an array of size `N` that indicates which of the `N` largest integers smaller than `upper_limit` are prime.
@@ -152,12 +171,16 @@ pub const fn sieve_lt<const N: usize, const MEM: usize>(
     let base_sieve: [bool; MEM] = sieve();
 
     // Use the result to sieve the higher range.
-    let upper_sieve = sieve_segment(&base_sieve, upper_limit);
+    let (offset, upper_sieve) = match sieve_segment(&base_sieve, upper_limit) {
+        Ok(res) => (0, res),
+        // The sieve contained more entries than there are non-negative numbers below the upper limit.
+        Err(_) => ((MEM as u64 - upper_limit) as usize, base_sieve),
+    };
 
-    let mut ans = [false; N];
     let mut i = 0;
+    let mut ans = [false; N];
     while i < N {
-        ans[N - 1 - i] = upper_sieve[MEM - 1 - i];
+        ans[N - 1 - i] = upper_sieve[MEM - 1 - i - offset];
         i += 1;
     }
     Ok(ans)
@@ -331,7 +354,10 @@ pub const fn sieve_geq<const N: usize, const MEM: usize>(
 
     let base_sieve: [bool; MEM] = sieve();
 
-    let upper_sieve = sieve_segment(&base_sieve, upper_limit);
+    let upper_sieve = match sieve_segment(&base_sieve, upper_limit) {
+        Ok(res) => res,
+        Err(_) => panic!("this is already checked above"),
+    };
 
     let mut ans = [false; N];
     let mut i = 0;
@@ -397,13 +423,25 @@ macro_rules! sieve_segment {
 
 #[cfg(test)]
 mod test {
+    use crate::sieve::SegmentedSieveError;
+
     use super::{sieve, sieve_segment};
 
     #[test]
     fn test_consistency_of_sieve_segment() {
-        const P: [bool; 10] = sieve_segment(&sieve(), 10);
-        const PP: [bool; 10] = sieve_segment(&sieve(), 11);
+        const P: [bool; 10] = match sieve_segment(&sieve(), 10) {
+            Ok(s) => s,
+            Err(_) => panic!(),
+        };
+        const PP: [bool; 10] = match sieve_segment(&sieve(), 11) {
+            Ok(s) => s,
+            Err(_) => panic!(),
+        };
         assert_eq!(P, sieve());
         assert_eq!(PP, sieve::<11>()[1..]);
+        assert_eq!(
+            sieve_segment::<5>(&[false, false, true, true, false], 4),
+            Err(SegmentedSieveError)
+        );
     }
 }