From cfd9ddb6757e0edb3112e33c722d591bd323104e Mon Sep 17 00:00:00 2001
From: Chip Hogg <chogg@aurora.tech>
Date: Wed, 13 Nov 2024 19:56:05 -0500
Subject: [PATCH] Add Miller-Rabin probable prime test

This can mark a number as either "probably prime", or "definitely
composite".  The first parameter is the base, and the second is the
number to test.

Future PRs will build up the Strong Lucas test which complements this,
and then form the Baillie-PSW test by composing the two.

Helps #506.
---
 src/core/include/mp-units/ext/prime.h | 45 +++++++++++++++++++++++++++
 test/static/prime_test.cpp            | 18 +++++++++++
 2 files changed, 63 insertions(+)

diff --git a/src/core/include/mp-units/ext/prime.h b/src/core/include/mp-units/ext/prime.h
index c734c3d3..51050f41 100644
--- a/src/core/include/mp-units/ext/prime.h
+++ b/src/core/include/mp-units/ext/prime.h
@@ -138,6 +138,51 @@ namespace mp_units::detail {
   return result;
 }
 
+// Decompose a positive integer by factoring out all powers of 2.
+struct NumberDecomposition {
+  std::uint64_t power_of_two;
+  std::uint64_t odd_remainder;
+};
+
+// Express any positive `n` as `(2^s * d)`, where `d` is odd.
+[[nodiscard]] consteval NumberDecomposition decompose(std::uint64_t n)
+{
+  NumberDecomposition result{0u, n};
+  while (result.odd_remainder % 2u == 0u) {
+    result.odd_remainder /= 2u;
+    ++result.power_of_two;
+  }
+  return result;
+}
+
+// Perform a Miller-Rabin primality test on `n` using base `a`.
+//
+// Precondition: (a >= 2).
+// Precondition: (n >= a + 2).
+// Precondition: (n % 2 == 1).
+[[nodiscard]] consteval bool miller_rabin_probable_prime(std::uint64_t a, std::uint64_t n)
+{
+  MP_UNITS_EXPECTS_DEBUG(a >= 2u);
+  MP_UNITS_EXPECTS_DEBUG(n >= a + 2u);
+  MP_UNITS_EXPECTS_DEBUG(n % 2u == 1u);
+
+  const auto [s, d] = decompose(n - 1u);
+  auto x = pow_mod(a, d, n);
+  if (x == 1u) {
+    return true;
+  }
+
+  const auto minus_one = n - 1u;
+  for (auto r = 0u; r < s; ++r) {
+    if (x == minus_one) {
+      return true;
+    }
+    x = mul_mod(x, x, n);
+  }
+
+  return false;
+}
+
 [[nodiscard]] consteval bool is_prime_by_trial_division(std::uintmax_t n)
 {
   for (std::uintmax_t f = 2; f * f <= n; f += 1 + (f % 2)) {
diff --git a/test/static/prime_test.cpp b/test/static/prime_test.cpp
index bf1fd922..bfd7aae5 100644
--- a/test/static/prime_test.cpp
+++ b/test/static/prime_test.cpp
@@ -104,4 +104,22 @@ static_assert(half_mod_odd(MAX_U64 - 2u, MAX_U64) == MAX_U64 - 1u);
 static_assert(pow_mod(5u, 8u, 9u) == ((5u * 5u * 5u * 5u) * (5u * 5u * 5u * 5u)) % 9u);
 static_assert(pow_mod(2u, 64u, MAX_U64) == 1u);
 
+// Miller-Rabin primality testing.
+static_assert(miller_rabin_probable_prime(2u, 5u));
+static_assert(miller_rabin_probable_prime(2u, 7u));
+static_assert(!miller_rabin_probable_prime(2u, 9u));
+static_assert(miller_rabin_probable_prime(2u, 11u));
+
+static_assert(miller_rabin_probable_prime(2u, 2047u), "Known base 2 pseudoprime");
+static_assert(miller_rabin_probable_prime(2u, 3277u), "Known base 2 pseudoprime");
+
+static_assert(miller_rabin_probable_prime(3u, 121u), "Known base 3 pseudoprime");
+static_assert(miller_rabin_probable_prime(3u, 703u), "Known base 3 pseudoprime");
+
+static_assert(miller_rabin_probable_prime(2u, 225'653'407'801u), "Large known prime");
+static_assert(miller_rabin_probable_prime(2u, 334'524'384'739u), "Large known prime");
+static_assert(miller_rabin_probable_prime(2u, 9'007'199'254'740'881u), "Large known prime");
+
+static_assert(miller_rabin_probable_prime(2u, 18'446'744'073'709'551'557u), "Largest 64-bit prime");
+
 }  // namespace