Got segmented sieve version working

src/main.rs

@@ -39,7 +39,7 @@ const fn ln(n: f64) -> f64 {
 }
 
 const fn prime_limit(n: usize) -> usize {
-    if n >= 6 {
+    if n > 6 {
         let ln_n = ln(n as f64);
         (n as f64 * (ln_n + ln(ln_n))).ceil() as usize
     } else {
@@ -47,31 +47,17 @@ const fn prime_limit(n: usize) -> usize {
     }
 }
 
-const fn sieve_nth_prime<const N: usize>(n: usize) -> Option<usize> {
+const fn simple_sieve<const L: usize>(limit: usize) -> ([usize; L], usize) {
-    if n < 1 {
+    let mut is_prime = [true; L];
-        return None;
+    is_prime[0] = false;
-    }
+    is_prime[1] = false;
-    if n == 1 {
-        return Some(2);
-    }
-
-    let limit = prime_limit(n);
-    let mut sieve = [true; N];
-    let mut primes = [0usize; N];
-    let mut prime_index: usize = 0;
-
-    let mut p: usize = 2;
-    while p <= limit {
-        if sieve[p] {
-            primes[prime_index] = p;
-            prime_index += 1;
-            if prime_index + 1 == n {
-                return Some(p);
-            }
 
+    let mut p = 2;
+    while p * p <= limit {
+        if is_prime[p] {
             let mut i = p * p;
             while i <= limit {
-                sieve[i] = false;
+                is_prime[i] = false;
                 i += p;
             }
         }
@@ -79,21 +65,91 @@ const fn sieve_nth_prime<const N: usize>(n: usize) -> Option<usize> {
         p += 1;
     }
 
+    let mut primes = [0; L];
+    let mut n = 2;
+    let mut c = 0;
+
+    while n < limit {
+        if is_prime[n] {
+            primes[c] = n;
+            c += 1;
+        }
+
+        n += 1;
+    }
+
+    (primes, c)
+}
+
+const fn segmented_sieve<const L: usize>(n: usize) -> Option<usize> {
+    if n < 1 {
+        return None;
+    }
+    if n <= 6 {
+        return Some([2, 3, 5, 7, 11, 13][n - 1]);
+    }
+
+    let upper = prime_limit(n);
+    let limit = upper.isqrt() + 1;
+    let (primes, mut count) = simple_sieve::<L>(limit);
+
+    let mut low = limit;
+    let mut high = limit * 2;
+
+    while low < upper {
+        if high >= upper {
+            high = upper;
+        }
+
+        let mut mark = [true; L];
+
+        let mut i = 0;
+        while primes[i] != 0 {
+            let mut lo_lim = (low / primes[i]) * primes[i];
+
+            if lo_lim < low {
+                lo_lim += primes[i];
+            }
+
+            let mut j = lo_lim;
+            while j < high {
+                mark[j - low] = false;
+                j += primes[i];
+            }
+
+            i += 1;
+        }
+
+        let mut k = low;
+        while k < high {
+            if mark[k - low] {
+                count += 1;
+                if count == n {
+                    return Some(k);
+                }
+            }
+            k += 1;
+        }
+
+        low = low + limit;
+        high = high + limit;
+    }
+
     unreachable!()
 }
 
 macro_rules! nth_prime {
     ($n:expr) => {
         const {
-            const LIMIT: usize = prime_limit($n + 1) + 1;
+            const LIMIT: usize = prime_limit($n).isqrt() + 2; // Space complexity: O(sqrt(n))
-            sieve_nth_prime::<LIMIT>($n + 1).unwrap()
+            segmented_sieve::<LIMIT>($n).unwrap()
         }
     };
 }
 
 #[allow(long_running_const_eval)]
-const NTH_PRIME: usize = nth_prime!(1000000);
+const NTH_PRIME: usize = nth_prime!(1000);
 
 fn main() {
-    println!("{NTH_PRIME}")
+    println!("{NTH_PRIME}");
 }