Miller Rabin primality test

miller_rabin.py

#!/usr/bin/env python3

from random import randint

def miller_rabin(n: int) -> bool:
    if n <= 3 or n % 2 == 0:
        raise RuntimeError("n <= 3 or n % 2 == 0")

    k = 1
    while (n - 1) % (1 << k) == 0 and ((n - 1) // (1 << k)) % 2 == 0:
        k += 1

    m = (n - 1) // (1 << k)
    assert m % 2 != 0

    a = randint(2, n - 2)

    b = [pow(a, m, n)] + [0 for _ in range(k - 1)]
    assert len(b) == k

    if b[0] == 1 or b[0] == -1 % n:
        return True

    for i in range(1, k):
        b[i] = pow(b[i - 1], 2, n)
        if b[i] == 1:
            return False

        if b[i] == -1 % n:
            return True

    return False

def main() -> None:
    n = int(input("Enter integer: "))

    if miller_rabin(n):
        print("~PRIME")
    else:
        print("COMPOSITE")

if __name__ == "__main__":
    main()