Merge pull request #638 from chiphogg/chiphogg/miller-rabin#509

Add Miller-Rabin probable prime test
mpusz · Nov 14, 2024 · 7e1baf6 · 7e1baf6
2 parents e4bccad + cfd9ddb
commit 7e1baf6
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 0 deletions.
diff --git 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
@@ -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