Fix wheel factorization algorithm

src/core/include/units/bits/prime.h

@@ -26,6 +26,7 @@
 #include <cassert>
 #include <cstddef>
 #include <optional>
+#include <tuple>
 
 namespace units::detail
 {
@@ -37,32 +38,6 @@ constexpr bool is_prime_by_trial_division(std::size_t n) {
   return true;
 }
 
-constexpr std::size_t num_primes_between(std::size_t start, std::size_t end) {
-  std::size_t n = 0;
-
-  for (auto k = start; k < end; ++k) {
-    if (is_prime_by_trial_division(k)) { ++n; }
-  }
-
-  return n;
-}
-
-template<std::size_t Start, std::size_t End>
-constexpr auto primes_between() {
-  std::array<std::size_t, num_primes_between(Start, End)> results;
-  auto iter = std::begin(results);
-
-  for (auto k = Start; k < End; ++k) {
-    if (is_prime_by_trial_division(k)) {
-      *iter = k;
-      ++iter;
-    }
-  }
-
-  assert(iter == std::end(results));
-  return results;
-}
-
 // Return the first factor of n, as long as it is either k or n.
 //
 // Precondition: no integer smaller than k evenly divides n.
@@ -83,6 +58,34 @@ constexpr std::array<std::size_t, N> first_n_primes() {
   return primes;
 }
 
+template<std::size_t N, typename Callable>
+constexpr void call_for_coprimes_up_to(std::size_t n, const std::array<std::size_t, N> &basis, Callable &&call) {
+  for (std::size_t i = 0u; i < n; ++i) {
+    if (std::apply([&i](auto... primes) { return ((i % primes != 0) && ...); }, basis)) {
+      call(i);
+    }
+  }
+}
+
+template<std::size_t N>
+constexpr std::size_t num_coprimes_up_to(std::size_t n, const std::array<std::size_t, N> &basis) {
+  std::size_t count = 0u;
+  call_for_coprimes_up_to(n, basis, [&count](auto) { ++count; });
+  return count;
+}
+
+template<std::size_t ResultSize, std::size_t N>
+constexpr auto coprimes_up_to(std::size_t n, const std::array<std::size_t, N> &basis) {
+  std::array<std::size_t, ResultSize> coprimes;
+  std::size_t i = 0u;
+
+  call_for_coprimes_up_to(n, basis, [&coprimes, &i](std::size_t cp) { coprimes[i++] = cp; });
+
+  assert (i == N);
+
+  return coprimes;
+}
+
 template<std::size_t N>
 constexpr std::size_t product(const std::array<std::size_t, N> &values) {
   std::size_t product = 1;
@@ -93,41 +96,42 @@ constexpr std::size_t product(const std::array<std::size_t, N> &values) {
 // A configurable instantiation of the "wheel factorization" algorithm [1] for prime numbers.
 //
 // Instantiate with N to use a "basis" of the first N prime numbers.  Higher values of N use fewer trial divisions, at
-// the cost of additional space.  The amount of space consumed is roughly the total number of prime numbers that are
-// less than the _product_ of the first N prime numbers.  This means it grows rapidly with N.  Consider this approximate
-// chart:
+// the cost of additional space.  The amount of space consumed is roughly the total number of numbers that are a) less
+// than the _product_ of the basis elements (first N primes), and b) coprime with every element of the basis.  This
+// means it grows rapidly with N.  Consider this approximate chart:
 //
-//   N | Num primes stored | Trial divisions needed ~= (stored + 1 - N) / product(basis)
-//   --+-------------------+-----------------------
-//   1 |                 1 |                 50.0 %
-//   2 |                 3 |                 33.3 %
-//   3 |                10 |                 26.7 %
-//   4 |                46 |                 20.5 %
-//   5 |               343 |                 14.7 %
+//   N | Num coprimes | Trial divisions needed
+//   --+--------------+-----------------------
+//   1 |            1 |                 50.0 %
+//   2 |            2 |                 33.3 %
+//   3 |            8 |                 26.7 %
+//   4 |           48 |                 22.9 %
+//   5 |          480 |                 20.8 %
 //
-// Consider this behaviour when choosing the value of N most appropriate for your needs, and watch out for diminishing
-// returns.
+// Note the diminishing returns, and the rapidly escalating costs.  Consider this behaviour when choosing the value of N
+// most appropriate for your needs.
 //
 // [1] https://en.wikipedia.org/wiki/Wheel_factorization
 template<std::size_t BasisSize>
 struct WheelFactorizer {
   static constexpr auto basis = first_n_primes<BasisSize>();
   static constexpr std::size_t wheel_size = product(basis);
-  static constexpr auto primes_in_first_wheel = primes_between<basis.back() + 1, wheel_size>();
+  static constexpr auto coprimes_in_first_wheel =
+    coprimes_up_to<num_coprimes_up_to(wheel_size, basis)>(wheel_size, basis);
 
   static constexpr std::size_t find_first_factor(std::size_t n) {
     for (const auto& p : basis) {
       if (const auto k = first_factor_maybe(n, p)) { return *k; }
     }
 
-    for (const auto& p : primes_in_first_wheel) {
-      if (const auto k = first_factor_maybe(n, p)) { return *k; }
+    for (auto it = std::next(std::begin(coprimes_in_first_wheel)); it != std::end(coprimes_in_first_wheel); ++it) {
+      if (const auto k = first_factor_maybe(n, *it)) { return *k; }
     }
 
     for (std::size_t wheel = wheel_size; wheel < n; wheel += wheel_size) {
       if (const auto k = first_factor_maybe(n, wheel + 1)) { return *k; }
 
-      for (const auto &p : primes_in_first_wheel) {
+      for (const auto &p : coprimes_in_first_wheel) {
         if (const auto k = first_factor_maybe(n, wheel + p)) { return *k; }
       }
     }

test/unit_test/static/prime_test.cpp

@@ -35,17 +35,30 @@ constexpr bool check_primes(std::index_sequence<Is...>) {
 
 static_assert(check_primes<2>(std::make_index_sequence<122>{}));
 
-constexpr auto some_primes = primes_between<7, 19>(); // 7, 11, 13, 17
-static_assert(std::is_same_v<std::remove_cvref_t<decltype(some_primes)>, std::array<std::size_t, 4>>);
-static_assert(some_primes[0] == 7);
-static_assert(some_primes[1] == 11);
-static_assert(some_primes[2] == 13);
-static_assert(some_primes[3] == 17);
+// This is the smallest number that can catch the bug where we use only _prime_ numbers in the first wheel, rather than
+// numbers which are _coprime to the basis_.
+//
+// The basis for N = 4 is {2, 3, 5, 7}, so the wheel size is 210.  11 * 11 = 121 is within the first wheel.  It is
+// coprime with every element of the basis, but it is _not_ prime.  If we keep only prime numbers, then we will neglect
+// using numbers of the form (210 * n + 121) as trial divisors, which is a problem if any are prime.  For n = 1, we have
+// a divisor of (210 + 121 = 331), which happens to be prime but will not be used.  Thus, (331 * 331 = 109561) is a
+// composite number which could wrongly appear prime if we skip over 331.
+static_assert(WheelFactorizer<4>::is_prime(109561) == is_prime_by_trial_division(109561));
 
-static_assert(WheelFactorizer<1>::primes_in_first_wheel.size() == 0);
-static_assert(WheelFactorizer<2>::primes_in_first_wheel.size() == 1);
-static_assert(WheelFactorizer<3>::primes_in_first_wheel.size() == 7);
-static_assert(WheelFactorizer<4>::primes_in_first_wheel.size() == 42);
+static_assert(WheelFactorizer<1>::coprimes_in_first_wheel.size() == 1);
+static_assert(WheelFactorizer<2>::coprimes_in_first_wheel.size() == 2);
+static_assert(WheelFactorizer<3>::coprimes_in_first_wheel.size() == 8);
+static_assert(WheelFactorizer<4>::coprimes_in_first_wheel.size() == 48);
+static_assert(WheelFactorizer<5>::coprimes_in_first_wheel.size() == 480);
+
+static_assert(WheelFactorizer<3>::coprimes_in_first_wheel[0] == 1);
+static_assert(WheelFactorizer<3>::coprimes_in_first_wheel[1] == 7);
+static_assert(WheelFactorizer<3>::coprimes_in_first_wheel[2] == 11);
+static_assert(WheelFactorizer<3>::coprimes_in_first_wheel[3] == 13);
+static_assert(WheelFactorizer<3>::coprimes_in_first_wheel[4] == 17);
+static_assert(WheelFactorizer<3>::coprimes_in_first_wheel[5] == 19);
+static_assert(WheelFactorizer<3>::coprimes_in_first_wheel[6] == 23);
+static_assert(WheelFactorizer<3>::coprimes_in_first_wheel[7] == 29);
 
 static_assert(!WheelFactorizer<1>::is_prime(0));
 static_assert(!WheelFactorizer<1>::is_prime(1));