Implement Strong Lucas test and Baillie-PSW

au/code/au/utility/probable_primes.hh

@@ -84,6 +84,27 @@ constexpr PrimeResult miller_rabin(std::size_t a, uint64_t n) {
     return PrimeResult::COMPOSITE;
 }
 
+//
+// Test whether the number is a perfect square.
+//
+constexpr bool is_perfect_square(uint64_t n) {
+    if (n < 2u) {
+        return true;
+    }
+
+    uint64_t prev = n / 2u;
+    while (true) {
+        const uint64_t curr = (prev + n / prev) / 2u;
+        if (curr * curr == n) {
+            return true;
+        }
+        if (curr >= prev) {
+            return false;
+        }
+        prev = curr;
+    }
+}
+
 constexpr uint64_t gcd(uint64_t a, uint64_t b) {
     while (b != 0u) {
         const auto remainder = a % b;
@@ -164,5 +185,166 @@ constexpr int jacobi_symbol(int64_t raw_a, uint64_t n) {
     return jacobi_symbol_positive_numerator(a, n, result);
 }
 
+// The "D" parameter in the Strong Lucas probable prime test.
+//
+// Default construction produces the first value to try according to Selfridge's parameter
+// selection.  Calling `increment()` on this will successively produce the next parameter to try.
+struct LucasDParameter {
+    uint64_t mag = 5u;
+    bool is_positive = true;
+
+    friend constexpr int as_int(const LucasDParameter &D) {
+        return bool_sign(D.is_positive) * static_cast<int>(D.mag);
+    }
+    friend constexpr void increment(LucasDParameter &D) {
+        D.mag += 2u;
+        D.is_positive = !D.is_positive;
+    }
+};
+
+//
+// The first `D` in the infinite sequence {5, -7, 9, -11, ...} whose Jacobi symbol is (-1) is the
+// `D` we want to use for the Strong Lucas Probable Prime test.
+//
+// Requires that `n` is *not* a perfect square.
+//
+constexpr LucasDParameter find_first_D_with_jacobi_symbol_neg_one(uint64_t n) {
+    LucasDParameter D{};
+    while (jacobi_symbol(as_int(D), n) != -1) {
+        increment(D);
+    }
+    return D;
+}
+
+//
+// Elements of the Lucas sequence.
+//
+// The default values give the first element (i.e., k=1) of the sequence.
+//
+struct LucasSequenceElement {
+    uint64_t U = 1u;
+    uint64_t V = 1u;
+};
+
+// Produce the Lucas element whose index is twice the input element's index.
+constexpr LucasSequenceElement double_strong_lucas_index(const LucasSequenceElement &element,
+                                                         uint64_t n,
+                                                         LucasDParameter D) {
+    const auto &U = element.U;
+    const auto &V = element.V;
+
+    uint64_t V_squared = mul_mod(V, V, n);
+    uint64_t D_U_squared = mul_mod(D.mag, mul_mod(U, U, n), n);
+    uint64_t V2 =
+        D.is_positive ? add_mod(V_squared, D_U_squared, n) : sub_mod(V_squared, D_U_squared, n);
+    V2 = half_mod_odd(V2, n);
+
+    return LucasSequenceElement{
+        mul_mod(U, V, n),
+        V2,
+    };
+}
+
+// Find the next element in the Lucas sequence, using parameters for strong Lucas probable primes.
+constexpr LucasSequenceElement increment_strong_lucas_index(const LucasSequenceElement &element,
+                                                            uint64_t n,
+                                                            LucasDParameter D) {
+    const auto &U = element.U;
+    const auto &V = element.V;
+
+    auto U2 = half_mod_odd(add_mod(U, V, n), n);
+
+    const auto D_U = mul_mod(D.mag, U, n);
+    auto V2 = D.is_positive ? add_mod(V, D_U, n) : sub_mod(V, D_U, n);
+    V2 = half_mod_odd(V2, n);
+
+    return LucasSequenceElement{U2, V2};
+}
+
+// Compute the strong Lucas sequence element at index `i`.
+constexpr LucasSequenceElement find_strong_lucas_element(uint64_t i,
+                                                         uint64_t n,
+                                                         LucasDParameter D) {
+    LucasSequenceElement element{};
+
+    bool bits[64] = {};
+    std::size_t n_bits = 0u;
+    while (i > 1u) {
+        bits[n_bits++] = (i % 2u == 1u);
+        i /= 2u;
+    }
+
+    for (std::size_t j = n_bits; j > 0u; --j) {
+        element = double_strong_lucas_index(element, n, D);
+        if (bits[j - 1u]) {
+            element = increment_strong_lucas_index(element, n, D);
+        }
+    }
+
+    return element;
+}
+
+//
+// Perform a strong Lucas primality test on `n`.
+//
+constexpr PrimeResult strong_lucas(uint64_t n) {
+    if (n < 2u || n % 2u == 0u) {
+        return PrimeResult::BAD_INPUT;
+    }
+
+    if (is_perfect_square(n)) {
+        return PrimeResult::COMPOSITE;
+    }
+
+    const auto D = find_first_D_with_jacobi_symbol_neg_one(n);
+
+    const auto params = decompose(n + 1u);
+    const auto &s = params.power_of_two;
+    const auto &d = params.odd_remainder;
+
+    auto element = find_strong_lucas_element(d, n, D);
+    if (element.U == 0u) {
+        return PrimeResult::PROBABLY_PRIME;
+    }
+
+    for (std::size_t i = 0u; i < s; ++i) {
+        if (element.V == 0u) {
+            return PrimeResult::PROBABLY_PRIME;
+        }
+        element = double_strong_lucas_index(element, n, D);
+    }
+
+    return PrimeResult::COMPOSITE;
+}
+
+//
+// Perform the Baillie-PSW test for primality.
+//
+// Returns `BAD_INPUT` for any number less than 2, `COMPOSITE` for any larger number that is _known_
+// to be prime, and `PROBABLY_PRIME` for any larger number that is deemed "probably prime", which
+// includes all prime numbers.
+//
+// Actually, the Baillie-PSW test is known to be completely accurate for all 64-bit numbers;
+// therefore, since our input type is `uint64_t`, the output will be `PROBABLY_PRIME` if and only if
+// the input is prime.
+//
+constexpr PrimeResult baillie_psw(uint64_t n) {
+    if (n < 2u) {
+        return PrimeResult::BAD_INPUT;
+    }
+    if (n < 4u) {
+        return PrimeResult::PROBABLY_PRIME;
+    }
+    if (n % 2u == 0u) {
+        return PrimeResult::COMPOSITE;
+    }
+
+    if (miller_rabin(2u, n) == PrimeResult::COMPOSITE) {
+        return PrimeResult::COMPOSITE;
+    }
+
+    return strong_lucas(n);
+}
+
 }  // namespace detail
 }  // namespace au

au/code/au/utility/test/probable_primes_test.cc

@@ -200,6 +200,115 @@ TEST(MillerRabin, SupportsConstexpr) {
     static_assert(result == PrimeResult::PROBABLY_PRIME, "997 is prime");
 }
 
+TEST(IsPerfectSquare, ProducesCorrectAnswers) {
+    auto next_sqrt = 0u;
+    for (auto n = 0u; n < 400'000u; ++n) {
+        const auto next_square = next_sqrt * next_sqrt;
+
+        const auto is_square = (n == next_square);
+        if (is_square) {
+            ++next_sqrt;
+        }
+
+        EXPECT_EQ(is_perfect_square(n), is_square) << "n = " << n;
+    }
+}
+
+std::vector<uint64_t> strong_lucas_pseudoprimes() {
+    // https://oeis.org/A217255
+    return {5459u,   5777u,   10877u,  16109u,  18971u,  22499u,  24569u,  25199u,  40309u,
+            58519u,  75077u,  97439u,  100127u, 113573u, 115639u, 130139u, 155819u, 158399u,
+            161027u, 162133u, 176399u, 176471u, 189419u, 192509u, 197801u, 224369u, 230691u,
+            231703u, 243629u, 253259u, 268349u, 288919u, 313499u, 324899u};
+}
+
+TEST(LucasDParameter, CanConvertToInt) {
+    EXPECT_EQ(as_int(LucasDParameter{5u, true}), 5);
+    EXPECT_EQ(as_int(LucasDParameter{7u, false}), -7);
+}
+
+TEST(StrongLucas, AllPrimeNumbersAreProbablyPrime) {
+    const auto primes = first_n_primes<3'000u>();
+    for (const auto &p : primes) {
+        if (p > 2u) {
+            EXPECT_THAT(strong_lucas(p), Eq(PrimeResult::PROBABLY_PRIME)) << p;
+        }
+    }
+}
+
+TEST(StrongLucas, GetsFooledByKnownPseudoprimes) {
+    for (const auto &p : strong_lucas_pseudoprimes()) {
+        ASSERT_EQ(miller_rabin(2u, p), PrimeResult::COMPOSITE) << p;
+        EXPECT_THAT(strong_lucas(p), Eq(PrimeResult::PROBABLY_PRIME)) << p;
+    }
+}
+
+TEST(StrongLucas, OddNumberIsProbablyPrimeIffPrimeOrPseudoprime) {
+    const auto primes = first_n_primes<3'000u>();
+    const auto pseudoprimes = strong_lucas_pseudoprimes();
+
+    // Make sure that we are both _into the regime_ of the pseudoprimes, and that we aren't off the
+    // end of it.
+    ASSERT_THAT(primes.back(), AllOf(Gt(pseudoprimes.front()), Lt(pseudoprimes.back())));
+
+    std::size_t i_prime = 1u;  // Skip 2; we're only checking odd numbers.
+    std::size_t i_pseudoprime = 0u;
+    for (uint64_t n = primes[i_prime]; i_prime < primes.size(); n += 2u) {
+        const auto is_prime = (n == primes[i_prime]);
+        if (is_prime) {
+            ++i_prime;
+        }
+
+        const auto is_pseudoprime = (n == pseudoprimes[i_pseudoprime]);
+        if (is_pseudoprime) {
+            ++i_pseudoprime;
+        }
+
+        const auto expected =
+            (is_prime || is_pseudoprime) ? PrimeResult::PROBABLY_PRIME : PrimeResult::COMPOSITE;
+        EXPECT_THAT(strong_lucas(n), Eq(expected)) << "n = " << n;
+    }
+}
+
+TEST(BailliePSW, BadInputForLessThanTwo) {
+    EXPECT_THAT(baillie_psw(0u), Eq(PrimeResult::BAD_INPUT));
+    EXPECT_THAT(baillie_psw(1u), Eq(PrimeResult::BAD_INPUT));
+}
+
+TEST(BailliePSW, TwoIsPrime) { EXPECT_THAT(baillie_psw(2u), Eq(PrimeResult::PROBABLY_PRIME)); }
+
+TEST(BailliePSW, CorrectlyIdentifiesAllOddNumbersUpToTheFirstThousandPrimes) {
+    const auto first_10k_primes = first_n_primes<10'000u>();
+
+    std::size_t i_prime = 1u;  // Skip "prime 0" (a.k.a. "2").
+    for (uint64_t i = 3u; i_prime < first_10k_primes.size(); i += 2u) {
+        const bool is_prime = (i == first_10k_primes[i_prime]);
+        if (is_prime) {
+            ++i_prime;
+        }
+        const auto expected = is_prime ? PrimeResult::PROBABLY_PRIME : PrimeResult::COMPOSITE;
+        EXPECT_THAT(baillie_psw(i), Eq(expected)) << "i = " << i;
+    }
+}
+
+TEST(BailliePSW, IdentifiesPerfectSquareAsComposite) {
+    // (1093 ^ 2 = 1,194,649) is the smallest strong pseudoprime to base 2 that is a perfect square.
+    constexpr auto n = 1093u * 1093u;
+    ASSERT_THAT(miller_rabin(2u, n), Eq(PrimeResult::PROBABLY_PRIME));
+    EXPECT_THAT(baillie_psw(n), Eq(PrimeResult::COMPOSITE));
+}
+
+TEST(BailliePSW, HandlesVeryLargePrimes) {
+    for (const auto &p : {
+             225'653'407'801ull,
+             334'524'384'739ull,
+             9'007'199'254'740'881ull,
+             18'446'744'073'709'551'557ull,
+         }) {
+        EXPECT_THAT(baillie_psw(p), Eq(PrimeResult::PROBABLY_PRIME)) << p;
+    }
+}
+
 TEST(Gcd, ResultIsAlwaysAFactorAndGCDFindsNoLargerFactor) {
     for (auto i = 0u; i < 500u; ++i) {
         for (auto j = 1u; j < i; ++j) {