From f1c6438b802006ab794c7542c48e71aa341dd4f4 Mon Sep 17 00:00:00 2001
From: Avery Winters <git@averywinters.org>
Date: Wed, 6 Nov 2024 12:59:36 -0600
Subject: [PATCH] Use much simpler and more readable trial division

---
 rust/nth-prime/src/lib.rs | 87 ++++-----------------------------------
 1 file changed, 9 insertions(+), 78 deletions(-)

diff --git a/rust/nth-prime/src/lib.rs b/rust/nth-prime/src/lib.rs
index c727a5a..dc46bb2 100644
--- a/rust/nth-prime/src/lib.rs
+++ b/rust/nth-prime/src/lib.rs
@@ -1,83 +1,14 @@
-pub fn nth(mut n: u32) -> u32 {
-    let upper_bound = upper_bound_nth_prime(n);
-    let upper_bound: usize = upper_bound as usize + 1;
-    let len = upper_bound.div_ceil(8);
-    let mut bytes = vec![u8::MAX; len];
-    let mask = (1 << 0) | (1 << 1);
-    bytes[0] &= !mask;
-    let bound: f64 = upper_bound as f64;
-    let bound = bound.sqrt();
-    let bound: usize = bound as usize;
-
-    for i in 2..=bound {
-        let mask = 1 << (i % 8);
-        let byte = bytes[i / 8];
-        if byte & mask != 0u8 {
-            let start = i * i;
-            for j in (start..=upper_bound).step_by(i.try_into().unwrap()) {
-                let mask = 1 << (j % 8);
-                bytes[j / 8] &= !mask;
-            }
-        }
-    }
-
-    for (i, byte) in bytes.iter().enumerate() {
-        let count = byte.count_ones();
-        if count >= n {
-            for j in 0..8 {
-                let mask = 1 << j;
-                if byte & mask != 0u8 {
-                    if n == 0 {
-                        return (8 * i + j) as u32;
-                    }
-                    n -= 1;
-                }
-            }
-        } else {
-            n -= count;
-        }
-    }
-    unreachable!()
+pub fn nth(n: u32) -> u32 {
+    (2..)
+        .filter(|&x| is_prime(x))
+        .nth(n as usize)
+        .expect("at least n 32-bit integers are prime")
 }
 
-fn upper_bound_nth_prime(n: u32) -> u32 {
-    // For small prime counts, the upper bound estimate is not big enough.
-    // Hardcode it to the 6th prime (13).
-    if n < 6 {
-        return 13;
-    }
-
-    let mut low = n;
-    let mut high = u32::MAX - 1;
-    while low < high {
-        let mid = low + (high - low) / 2;
-        let guess = {
-            let mid: f64 = mid.into();
-            mid / mid.ln()
-        };
-        if guess > n.into() {
-            high = mid - 1;
-        } else {
-            low = mid + 1;
-        }
-    }
-    high + 1
+fn is_prime(n: u32) -> bool {
+    (2..=isqrt(n)).all(|x| n % x != 0)
 }
 
-#[cfg(test)]
-mod test {
-    use crate::upper_bound_nth_prime;
-
-    #[test]
-    fn greater_first_20() {
-        let primes = [
-            2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71,
-        ];
-
-        for (i, &p) in primes.iter().enumerate() {
-            let i = i.try_into().unwrap();
-            let upper_bound = upper_bound_nth_prime(i);
-            assert!(upper_bound >= p, "{upper_bound} < {p}");
-        }
-    }
+fn isqrt(n: u32) -> u32 {
+    (n as f64).sqrt() as u32
 }