Miller-Rabin primality test (Python)

From LiteratePrograms

Other implementations: C | C, GMP | Clojure | Groovy | Java | Python | Ruby | Scala

The Miller-Rabin primality test is a simple probabilistic algorithm for determining whether a number is prime or composite that is easy to implement. It proves compositeness of a number using the following formulas:

Suppose 0 < a < n is coprime to n (this is easy to test using the GCD). Write the number n−1 as , where d is odd. Then, provided that all of the following formulas hold, n is composite:

for all

If a is chosen uniformly at random and n is prime, these formulas hold with probability 1/4. Thus, repeating the test for k random choices of a gives a probability of 1 − 1 / 4^k that the number is prime. Moreover, Gerhard Jaeschke showed that any 32-bit number can be deterministically tested for primality by trying only a=2, 7, and 61.

We will implement the test for arbitrary-precision integers, with Python's long type, which is automatically used when integers get big enough, so we don't have to do anything special.

When performing the Miller-Rabin pass, only a, n, d, and powers of a are large numbers; s and i are logarithmically smaller. We precompute the large integers 1 and n−1 for use in comparisons and squaring. We've replaced the squaring using modular_exponent with a modular multiply of a_to_power by itself:

<<Miller-Rabin pass>>=
def miller_rabin_pass(a, s, d, n):
    a_to_power = pow(a, d, n)
    if a_to_power == 1:
        return True
    for i in xrange(s-1):
        if a_to_power == n - 1:
            return True
        a_to_power = (a_to_power * a_to_power) % n
    return a_to_power == n - 1

Since our numbers are now arbitrarily large, we can no longer take advantage of Jaeschke's result to produce a deterministic implementation, but we can produce a highly accurate probabilistic implementation by simply running the test enough times to make 1/4^k very small, say 20 times.

<<Miller-Rabin>>=
def miller_rabin(n):
    compute s and d
    for repeat in xrange(20):
        a = 0
        while a == 0:
            a = random.randrange(n)
        if not miller_rabin_pass(a, s, d, n):
            return False
    return True

Finally, we use bit shifts to compute d rapidly, and test just its least significant bit to determine if it is odd:

<<compute s and d>>=
d = n - 1
s = 0
while d % 2 == 0:
    d >>= 1
    s += 1

For convenience we ignore the special cases where n ≤ 2.

Finally, some test code:

<<MillerRabin.py>>=
import random, sys
Miller-Rabin pass
Miller-Rabin
if __name__ == "__main__":
    if sys.argv[1] == "test":
        n = long(sys.argv[2])
        print (miller_rabin(n) and "PRIME" or "COMPOSITE")
    elif sys.argv[1] == "genprime":
        nbits = int(sys.argv[2])
        while True:
            p = random.getrandbits(nbits)
            condition possible prime
            if miller_rabin(p):
                print p
                break

We've augmented main() with the ability to generate a prime number of a specified number of bits. To generate the random number we'll use random.getrandbits() since this method is faster than random.randrange() for generating numbers given a desired length in bits. It does this by randomly selecting numbers of that size until it finds a prime one. We could use this to rapidly generate large primes for use in cryptography. Before we test for primality we should condition the number by setting the most significant bit and the least significant bit using bitwise or. Doing this ensures that the number is odd (all primes > 2 are odd) and that the number is exactly the number of bits we specified:

<<condition possible prime>>=
p |= 2**nbits | 1

Here's some sample output:

$ python MillerRabin.py test 516119616549881
PRIME
$ python MillerRabin.py test 516119616549887
COMPOSITE
$ python MillerRabin.py genprime 128
277003545840181465186202237665148044033
$ python MillerRabin.py genprime 512
69070670148391210520342329574373473828883687815325200193110546590937028225649711\
62269141121965617545613897865904269544243721348093953041840506282278098321

Download code