Use Baillie-PSW for prime factorization

au/code/au/utility/factoring.hh

@@ -16,18 +16,61 @@
 
 #include <cstdint>
 
+#include "au/utility/probable_primes.hh"
+
 namespace au {
 namespace detail {
 
+// Check whether a number is prime.
+constexpr bool is_prime(std::uintmax_t n) {
+    static_assert(sizeof(std::uintmax_t) <= sizeof(std::uint64_t),
+                  "Baillie-PSW only strictly guaranteed for 64-bit numbers");
+
+    return baillie_psw(n) == PrimeResult::PROBABLY_PRIME;
+}
+
+template <typename T = void>
+struct FirstPrimesImpl {
+    static constexpr uint16_t values[] = {
+        2,   3,   5,   7,   11,  13,  17,  19,  23,  29,  31,  37,  41,  43,  47,  53,  59,
+        61,  67,  71,  73,  79,  83,  89,  97,  101, 103, 107, 109, 113, 127, 131, 137, 139,
+        149, 151, 157, 163, 167, 173, 179, 181, 191, 193, 197, 199, 211, 223, 227, 229, 233,
+        239, 241, 251, 257, 263, 269, 271, 277, 281, 283, 293, 307, 311, 313, 317, 331, 337,
+        347, 349, 353, 359, 367, 373, 379, 383, 389, 397, 401, 409, 419, 421, 431, 433, 439,
+        443, 449, 457, 461, 463, 467, 479, 487, 491, 499, 503, 509, 521, 523, 541};
+    static constexpr std::size_t N = sizeof(values) / sizeof(values[0]);
+};
+template <typename T>
+constexpr uint16_t FirstPrimesImpl<T>::values[];
+template <typename T>
+constexpr std::size_t FirstPrimesImpl<T>::N;
+using FirstPrimes = FirstPrimesImpl<>;
+
 // Find the smallest factor which divides n.
 //
 // Undefined unless (n > 1).
 constexpr std::uintmax_t find_first_factor(std::uintmax_t n) {
-    if (n % 2u == 0u) {
-        return 2u;
+    const auto &first_primes = FirstPrimes::values;
+    const auto &n_primes = FirstPrimes::N;
+
+    // First, do trial division against the first N primes.
+    for (const auto &p : first_primes) {
+        if (n % p == 0u) {
+            return p;
+        }
+
+        if (p * p > n) {
+            return n;
+        }
     }
 
-    std::uintmax_t factor = 3u;
+    // If we got this far, and haven't found a factor nor terminated, do a fast primality check.
+    if (is_prime(n)) {
+        return n;
+    }
+
+    // If we're here, we know `n` is composite, so continue with trial division for all odd numbers.
+    std::uintmax_t factor = first_primes[n_primes - 1u] + 2u;
     while (factor * factor <= n) {
         if (n % factor == 0u) {
             return factor;
@@ -38,9 +81,6 @@ constexpr std::uintmax_t find_first_factor(std::uintmax_t n) {
     return n;
 }
 
-// Check whether a number is prime.
-constexpr bool is_prime(std::uintmax_t n) { return (n > 1) && (find_first_factor(n) == n); }
-
 // Find the largest power of `factor` which divides `n`.
 //
 // Undefined unless n > 0, and factor > 1.

au/code/au/utility/probable_primes.hh

@@ -179,7 +179,7 @@ constexpr int jacobi_symbol(int64_t raw_a, uint64_t n) {
     // Starting conditions: transform `a` to strictly non-negative values, setting `result` to the
     // sign we pick up from this operation (if any).
     int result = bool_sign((raw_a >= 0) || (n % 4u == 1u));
-    auto a = static_cast<uint64_t>(std::abs(raw_a)) % n;
+    auto a = static_cast<uint64_t>(raw_a * bool_sign(raw_a >= 0)) % n;
 
     // Delegate to an implementation which can only handle positive numbers.
     return jacobi_symbol_positive_numerator(a, n, result);

au/code/au/utility/test/factoring_test.cc

@@ -14,14 +14,32 @@
 
 #include "au/utility/factoring.hh"
 
+#include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
+using ::testing::Gt;
+using ::testing::Le;
+
 namespace au {
 namespace detail {
 namespace {
 std::uintmax_t cube(std::uintmax_t n) { return n * n * n; }
 }  // namespace
 
+TEST(FirstPrimes, HasOnlyPrimesInOrderAndDoesntSkipAny) {
+    const auto &first_primes = FirstPrimes::values;
+    const auto &n_primes = FirstPrimes::N;
+    auto i_prime = 0u;
+    for (auto i = 2u; i <= first_primes[n_primes - 1u]; ++i) {
+        if (i == first_primes[i_prime]) {
+            EXPECT_TRUE(is_prime(i)) << i;
+            ++i_prime;
+        } else {
+            EXPECT_FALSE(is_prime(i)) << i;
+        }
+    }
+}
+
 TEST(FindFirstFactor, ReturnsInputForPrimes) {
     EXPECT_EQ(find_first_factor(2u), 2u);
     EXPECT_EQ(find_first_factor(3u), 3u);
@@ -37,6 +55,22 @@ TEST(FindFirstFactor, FindsFirstFactor) {
     EXPECT_EQ(find_first_factor(cube(196961u)), 196961u);
 }
 
+TEST(FindFirstFactor, CanFactorNumbersWithLargePrimeFactor) {
+    // Small prime factors.
+    EXPECT_EQ(find_first_factor(2u * 9'007'199'254'740'881u), 2u);
+    EXPECT_EQ(find_first_factor(3u * 9'007'199'254'740'881u), 3u);
+
+    constexpr auto LAST_TRIAL_PRIME = FirstPrimes::values[FirstPrimes::N - 1u];
+
+    // Large prime factor from trial division.
+    ASSERT_THAT(541u, Le(LAST_TRIAL_PRIME));
+    EXPECT_EQ(find_first_factor(541u * 9'007'199'254'740'881u), 541u);
+
+    // Large prime factor higher than what we use for trial division.
+    ASSERT_THAT(1999u, Gt(LAST_TRIAL_PRIME));
+    EXPECT_EQ(find_first_factor(1999u * 9'007'199'254'740'881u), 1999u);
+}
+
 TEST(IsPrime, FalseForLessThan2) {
     EXPECT_FALSE(is_prime(0u));
     EXPECT_FALSE(is_prime(1u));
@@ -60,6 +94,17 @@ TEST(IsPrime, FindsPrimes) {
     EXPECT_FALSE(is_prime(196962u));
 }
 
+TEST(IsPrime, CanHandleVeryLargePrimes) {
+    for (const auto &p : {
+             uint64_t{225'653'407'801u},
+             uint64_t{334'524'384'739u},
+             uint64_t{9'007'199'254'740'881u},
+             uint64_t{18'446'744'073'709'551'557u},
+         }) {
+        EXPECT_TRUE(is_prime(p)) << p;
+    }
+}
+
 TEST(Multiplicity, CountsFactors) {
     constexpr std::uintmax_t n = (2u * 2u * 2u) * (3u) * (5u * 5u);
     EXPECT_EQ(multiplicity(2u, n), 3u);