From 6e0ea7976e838b06f130a2a2d116e797a073ce54 Mon Sep 17 00:00:00 2001
From: Konstantin Avadov <konstantin.avadov@yahoo.com>
Date: Thu, 30 Nov 2023 01:31:19 +0300
Subject: [PATCH] Implementation of wheel factorization.

Also, it fixes the bug of skipping the first two numbers after n. For example, in the previous version:

>>> galois.next_prime(100000034)
100000037

>>> galois.next_prime(100000035)
100000039

>>> galois.next_prime(100000036)
100000039

>>> galois.next_prime(100000037)
100000049
---
 src/galois/_prime.py | 46 ++++++++++++++++++++++++++++++--------------
 1 file changed, 32 insertions(+), 14 deletions(-)

diff --git a/src/galois/_prime.py b/src/galois/_prime.py
index c6f0434d1..7389046fb 100644
--- a/src/galois/_prime.py
+++ b/src/galois/_prime.py
@@ -170,14 +170,23 @@ def prev_prime(n: int) -> int:
     if n <= MAX_N:
         return PRIMES[bisect.bisect_right(PRIMES, n) - 1]
 
-    # TODO: Make this faster using wheel factorization
-    n = n - 1 if n % 2 == 0 else n  # The next possible prime (which is odd)
-    while True:
-        n -= 2  # Only check odds
-        if is_prime(n):
-            break
+    shifts = [29, 23, 19, 17, 13, 11, 7, 1] # Factorization wheel for basis {2, 3, 5}
+    base = n // 30 * 30 # Wheel factorization starting point
+    found = False # Success flag
+
+    while not found:
+        for shift in shifts:
+            i = base + shift # May be bigger than n
+
+            if i >= n:
+                continue
+
+            if is_prime(i):
+                found = True
+                break
+        base -= 30
 
-    return n
+    return i
 
 
 @export
@@ -210,14 +219,23 @@ def next_prime(n: int) -> int:
     if n < PRIMES[-1]:
         return PRIMES[bisect.bisect_right(PRIMES, n)]
 
-    # TODO: Make this faster using wheel factorization
-    n = n + 1 if n % 2 == 0 else n + 2  # The next possible prime (which is odd)
-    while True:
-        n += 2  # Only check odds
-        if is_prime(n):
-            break
+    shifts = [1, 7, 11, 13, 17, 19, 23, 29] # Factorization wheel for basis {2, 3, 5}
+    base = n // 30 * 30 # Wheel factorization starting point. May be less than n.
+    found = False # Success flag
+
+    while not found:
+        for shift in shifts:
+            i = base + shift
+
+            if i <= n:
+                continue
+
+            if is_prime(i):
+                found = True
+                break
+        base += 30
 
-    return n
+    return i
 
 
 @export