From 980bfe71a5d263d96fbdca3c4585e7ee044f316a Mon Sep 17 00:00:00 2001
From: Matt Borland <matt@mattborland.com>
Date: Sun, 4 Oct 2020 10:28:08 -0500
Subject: [PATCH] Fixes for multiprecision types [CI SKIP] Primes up to sqrt
 upper_bound are now captured as found

---
 .../detail/linear_prime_sieve.hpp             | 30 ++++++++++++-------
 test/test_prime_sieve.cpp                     |  5 ++++
 2 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/include/boost/math/special_functions/detail/linear_prime_sieve.hpp b/include/boost/math/special_functions/detail/linear_prime_sieve.hpp
index 1e1061a5f..b48918a75 100644
--- a/include/boost/math/special_functions/detail/linear_prime_sieve.hpp
+++ b/include/boost/math/special_functions/detail/linear_prime_sieve.hpp
@@ -56,7 +56,7 @@ void mark_sieve(RandomAccessIterator first, RandomAccessIterator last, Integer f
 }
 
 template<typename Bitset, typename Integer>
-inline void mark_sieve(Bitset& bits, Integer factor)
+inline void mark_sieve(Bitset& bits, const Integer factor)
 {
     for(Integer i {factor * factor}; i < bits.size(); i += factor)
     {
@@ -101,35 +101,43 @@ inline decltype(auto) stepanov_sieve(Integer upper_bound, OutputIterator resulta
 // TODO(mborland): Pass in execution policy. mark_sieve can readily be converted to std::for_each, but dynamic_bitset would need replaced with something
 //                 that has iterators
 template<typename Integer, typename OutputIterator>
-decltype(auto) wheel_sieve_of_eratosthenes(Integer upper_bound, OutputIterator resultant_primes)
+decltype(auto) wheel_sieve_of_eratosthenes(const Integer upper_bound, OutputIterator resultant_primes)
 {
+    if(upper_bound == 2)
+    {
+        *resultant_primes++ = static_cast<Integer>(2);
+        return resultant_primes;
+    }
+
     const Integer sqrt_upper_bound {static_cast<Integer>(std::floor(std::sqrt(static_cast<double>(upper_bound)))) + 1};
     boost::dynamic_bitset<> trial(static_cast<std::size_t>(upper_bound));
     trial.set();
     std::array<Integer, 3> primes {2, 3, 5}; // Wheel basis
     std::array<Integer, 8> wheel {7, 11, 13, 17, 19, 23, 29, 31}; // MOD 30 wheel
 
-    for(auto& prime : primes)
+    for(std::size_t i {}; i < primes.size(); ++i)
     {
-        mark_sieve(trial, prime);
-        *resultant_primes++ = prime;
+        mark_sieve(trial, primes[i]);
+        *resultant_primes++ = primes[i];
     }
 
-    for(auto& co_prime : wheel)
+    // Last value in the wheel is the starting point for the next step
+    for(std::size_t i {}; i < wheel.size() - 1; ++i)
     {
-        mark_sieve(trial, co_prime);
-        *resultant_primes++ = co_prime;
+        mark_sieve(trial, wheel[i]);
+        *resultant_primes++ = wheel[i];
     }
 
-    for(Integer i {wheel.back()}; i < static_cast<std::size_t>(sqrt_upper_bound); i += wheel.back() - 1)
+    for(Integer i {wheel.back()}; i < sqrt_upper_bound; i += wheel.back() - 1)
     {
         if(trial[static_cast<std::size_t>(i)] == true)
         {
-            mark_sieve(trial, i * i);
+            mark_sieve(trial, i);
+            *resultant_primes++ = i;
         }
     }
 
-    for(Integer i {wheel.back() + 2}; i < upper_bound; i += 2)
+    for(Integer i {sqrt_upper_bound % 2 == 0 ? sqrt_upper_bound + 1 : sqrt_upper_bound + 2}; i < upper_bound; i += 2)
     {
         if(trial[static_cast<std::size_t>(i)])
         {
diff --git a/test/test_prime_sieve.cpp b/test/test_prime_sieve.cpp
index 14a4ca81e..0363d0615 100644
--- a/test/test_prime_sieve.cpp
+++ b/test/test_prime_sieve.cpp
@@ -403,6 +403,11 @@ int main()
     test_stepanov_sieve<int>();
 
     test_wheel_sieve_of_eratosthenes<int>();
+    test_wheel_sieve_of_eratosthenes<int32_t>();
+    test_wheel_sieve_of_eratosthenes<int64_t>();
+    test_wheel_sieve_of_eratosthenes<uint32_t>();
+    test_wheel_sieve_of_eratosthenes<boost::multiprecision::cpp_int>();
+    test_wheel_sieve_of_eratosthenes<boost::multiprecision::mpz_int>();
 
     // Composite
     test_prime_sieve<int>();