Verify Miller-Rabin primality test

miller-rabin-prime.cc

#include <iostream>
#include <cstdint>
#include <array>
#include <ranges>
#include <cassert>
#include <bitset>

uint32_t pow_n(uint32_t a, uint32_t d, uint32_t n) {
    if (d == 0) return 1;
    if (d == 1) return a;
    uint32_t t = pow_n(a, d / 2, n);
    t = ((uint64_t)t * t) % n;
    if (d % 2 == 0) {
        return t;
    } else {
        return ((uint64_t)t * a) % n;
    }
};

bool test(uint32_t n, unsigned s, uint32_t d, uint32_t a) {
    //std::cout << "test(n = " << n << ", s = " << s << ", d = " << d << ", a = " << a << ")\n";
    uint32_t x = pow_n(a ,d ,n);
    //std::cout << "  x = " << x << std::endl;
    if (x == 1 || x == n - 1) return true;
    for (unsigned i = 1; i < s; ++i) {
        x = ((uint64_t)x * x) % n;
        if (x == n - 1) return true;
    }
    return false;
}

bool is_prime(uint32_t n) {
    static const std::array witnesses{2u, 3u, 5u, 7u, 11u, 13u, 17u, 19u, 23u, 29u, 31u, 37u};
    static const std::array bounds{
        2'047llu, 1'373'653llu, 25'326'001llu, 3'215'031'751llu,
        2'152'302'898'747llu, 3'474'749'660'383llu,
        341'550'071'728'321llu, 341'550'071'728'321llu /* no bounds for 19 */, 
        3'825'123'056'546'413'051llu,
        3'825'123'056'546'413'051llu /* no bound for 29 */,
        3'825'123'056'546'413'051llu /* no bound for 31 */,
        (unsigned long long)UINT64_MAX /* off by a bit but it's the last bounds */,
    };
    static_assert(witnesses.size() == bounds.size());

    if (n == 2) return true; // 2 is prime
    if (n % 2 == 0) return false; // other even numbers are not
    if (n <= witnesses.back()) { // I know the first few primes
        return (std::ranges::find(witnesses, n) != std::end(witnesses));
    }
    // write n = 2^s * d + 1 with d odd
    unsigned s = 0;
    uint32_t d = n - 1;
    while (d % 2 == 0) {
        ++s;
        d /= 2;
    }
    // test widtnesses until the bounds say it's a sure thing
    auto it = bounds.cbegin();
    for (auto a : witnesses) {
        //std::cout << a << " ";
        if (!test(n, s, d, a)) return false;
        if (n < *it++) return true;
    }
    return true;
}

template<std::size_t N>
auto composite() {
    std::bitset<N / 2 + 1> is_composite;
    for (uint32_t i = 3; (uint64_t)i * i < N; i += 2) {
        if (is_composite[i / 2]) continue;
        for (uint64_t j = i * i; j < N; j += 2 * i) is_composite[j / 2] = true;
    }

    return is_composite;
}

bool slow_prime(uint32_t n) {
    static const auto is_composite = composite<UINT32_MAX + 1llu>();
    if (n < 2) return false;
    if (n == 2) return true;
    if (n % 2 == 0) return false;
    return !is_composite.test(n / 2);
}

int main() {
    /*
    std::cout << "2047: ";
    bool fast = is_prime(2047);
    bool slow = slow_prime(2047);
    std::cout << (fast ? "fast prime" : "");
    std::cout << (slow ? "slow prime" : "");
    std::cout << std::endl;
    */
    //std::cout << "2: prime\n";
    for (uint64_t i = 0; i <= UINT32_MAX; ++i) {
        if (i % 1000000 == 1) { std::cout << "\r" << i << " "; std::cout.flush(); }
        bool fast = is_prime(i);
        bool slow = slow_prime(i);
        if (fast != slow) std::cout << i << std::endl;
        assert(fast == slow);
        //std::cout << i << ": " << (is_prime(i) ? "prime" : "") << std::endl;
    }
}