Miller-Rabin Primality Test in Python

Question

I have the following implementation of the Miller-Rabin test in Python:

from random import randint
import decimal
decimal.getcontext().prec = 10000000

def miller_rabin(n, k):
  if n == 2 or n == 3:
    return True
  elif n == 1 or n <= 0 or str(n).isdigit() == False or n % 2 == 0:
    return False
  else:
    s = 1
    while True:
      if (n - 1) % (2 ** (s + 1)) == 0:
        s += 1
      else:
        break
    d = (n - 1) / (2 ** s)
    for _ in range(k):
      a = randint(2, n - 2)
      x = (decimal.Decimal(a) ** decimal.Decimal(d)) % decimal.Decimal(n)
      for _ in range(s):
        y = x * x % n
        if y == 1 and x != 1 and x != n - 1:
          return False
        x = y
      if y != 1:
        return False
  return True

• n is the number to be tested,
• k is the number of rounds of testing to perform,
• s and d are used to write n in the form \$2^s*d\$, and
• a is a base such that \$2 \le a \le n−2\$.

decimal is there to accommodate some of the resulting decimal values that may appear.

It works fine from what I've tested thus far. However, I'm pretty sure it's not very efficient. How can I optimize my code?

Answer 1

Bug: decimal exponentiation (yes, decimal exponentiation is a bug)

This is wrong:

x = (decimal.Decimal(a) ** decimal.Decimal(d)) % decimal.Decimal(n)

Using very big decimals (prec = 10000000) may make this appear to work for some values. But there is no fixed precision that is sufficient for any input, for any fixed precision you can find an input such that a ** d is so large that it exceeds that precision. Apart from sometimes breaking, it's also always a performance bug. This line should be:

x = pow(a, d, n)

This modular exponentiation does not need huge precision, it only works with values that are about as big as n. So it should be efficient.

If it seems like d could be a float, well, it can't.

d = (n - 1) / (2 ** s)

The division should be integer division, or bitwise shift. 2^s by definition divides n - 1: the value of s is defined to be the highest integer such that 2^s divides n - 1 (the point is to divide the powers of 2 out of n - 1).

x * x % n

This could also be pow(x, 2, n) but here it's not critical.

Answer 2

Use the built-in is_integer instead of string coercion to check for integrality.

n == 1 or n <= 0 is just n <= 1.

Any time there's a return, you don't need to follow it with an else or elif.

It's important that you accept an optional random generator argument, for the purposes of reproducibility in unit tests. Also: add unit tests.

Never indent at two spaces; use 4.

It's unsurprising to see single letter variables in a math routine, but you should comment them all. (You still haven't offered explanations for x and y).

When factoring 2, you needn't exponentiate.

Don't use Decimal; Python has arbitrary-precision integers by default and as in the comments, Miller-Rabin is integral throughout.

Suggested

import random
from typing import Callable


def factor_2(n: int) -> tuple[int, int]:
    """
    Factor powers of 2 from n-1
    :param n: Mill-Rabin test subject; n must be odd.
    :return: s, d such that n-1 = 2**s * d (for odd d)

    >>> factor_2(5)
    (2, 1)

    >>> s, d = factor_2(67)
    >>> 2**s * d
    66
    >>> d % 2
    1

    >>> s, d = factor_2(1201)
    >>> 2**s * d
    1200
    >>> d % 2
    1
    """

    s = 0
    d = n - 1
    while d % 2 == 0:
        s += 1
        d //= 2
    return s, d


def miller_rabin(
    n: float | int,
    n_rounds: int,
    randint: Callable[[int, int], int] = random.randint,
) -> bool:
    """
    Stochastic primality test.
    :param n: Number to be tested
    :param n_rounds: Maximum number of rounds to run before assuming prime.
                     Called 'k' in some literature.
    :param randint: Optional, alternative randint for reproducibility
    :return: True if probably prime, false if definitely composite
    """
    if n == 2 or n == 3:  # Fast path for small primes
        return True
    if n <= 1 or n % 2 == 0:  # Fast path for small or even composites
        return False
    if isinstance(n, float):
        if n.is_integer():
            n = int(n)  # n is a whole float; coerce to int and continue
        else:
            return False  # n is non-integral, so composite by definition

    s, d = factor_2(n)

    for _ in range(n_rounds):
        # base such that 2 <= a <= n − 2
        a = randint(2, n - 2)
        x = pow(base=a, exp=d, mod=n)

        for _ in range(s):
            y = x*x % n
            if y == 1 and x != 1 and x != n - 1:
                return False
            x = y
        if y != 1:
            return False
    return True


def test() -> None:
    rand = random.Random(x=0)

    # Test first 20 numbers for primality
    for n, should_be_prime in enumerate((
        # 0 1 2 3 4  5 6 7 8 9
          0,0,1,1,0, 1,0,1,0,0,
        # 1011121314 1516171819
          0,1,0,1,0, 0,0,1,0,1,
    )):
        actual = miller_rabin(n=n, n_rounds=10, randint=rand.randint)
        assert actual == should_be_prime

    # Test large prime and composite, int and float forms
    assert miller_rabin(n=5264191, n_rounds=10, randint=rand.randint)
    assert miller_rabin(n=5264191., n_rounds=10, randint=rand.randint)
    assert not miller_rabin(n=5264193, n_rounds=10, randint=rand.randint)
    assert not miller_rabin(n=5264193., n_rounds=10, randint=rand.randint)
    assert not miller_rabin(n=5264193.1, n_rounds=10, randint=rand.randint)

    # Test to see that the function runs OK with the default rand implementation
    assert miller_rabin(n=2, n_rounds=1)
    assert isinstance(miller_rabin(n=2, n_rounds=3), bool)


if __name__ == '__main__':
    test()