Miller-Rabin and Lucas-Lehmer Primality Test

prime.py

import random
import math
import sys


def SimplePrime(n):
	"""
	Simple trial-division primality check.
	Returns True if n is prime, False otherwise.

	Method:
	- Reject even numbers > 2.
	- Check divisibility by all odd integers from 3 to sqrt(n).
	"""
	lim = int(math.ceil(n ** 0.5)) + 1  # Upper bound for trial division
	if n == 1:
		return False
	if n == 2:
		return True
	if n % 2 == 0:
		return False
	for i in range(3, lim, 2):  # Test only odd divisors
		if n % i == 0:
			return False  # Found a divisor, n is composite
	return True


def CheckStrongProbablePrimeToBaseA(a, d, s, n):
	"""
	Checks if n passes the strong probable prime (SPRP) test to base a.
	Naive implementation — recalculates a^(2^r * d) for each r.

	Parameters:
	- a: base to test against
	- d, s: from n-1 = (2^s) * d decomposition
	- n: the number being tested

	Returns:
	True if n passes SPRP test for base a, otherwise False.
	"""
	x = a ** d % n # Bottleneck
	if x == 1:  # First possible condition of SPRP test
		return True

	for r in range(0, s):
		x = a ** ((2 ** r) * d) % n
		if x == n - 1:  # Second possible condition: -1 mod n
			return True

	return False


def CheckStrongProbablePrimeToBaseA2(a, d, s, n):
	"""
	Optimized version of CheckStrongProbablePrimeToBaseA.
	Uses repeated squaring instead of recomputing powers from scratch.

	Parameters are the same as above.
	"""
	x = 1
	if d > 0:
		x = a % n
	x = pow(x, d, n) # From https://dusty.phillips.codes/2018/09/13/an-intermediate-guide-to-rsa/
	if x == 1 or x == n - 1:
		return True

	# Square repeatedly to check for n-1
	for r in range(1, s):
		x = pow(x, 2, n)
		if x == n - 1:
			return True

	return False


def MillerRabinPrimalityTest(n, checks=10) -> bool:
	"""
	Performs the Miller-Rabin primality test.

	Parameters:
	- n: number to test
	- checks: 
		* if int, number of random bases to use
		* if iterable, treated as a fixed set of bases

	Returns:
	True if n is probably prime, False if composite.

	Notes:
	- Deterministic for certain fixed base sets up to given bounds.
	- Probabilistic otherwise.
	"""
	if n == 1:
		return False
	if n == 2:
		return True
	if n % 2 == 0:
		return False  # Even numbers > 2 are not prime

	# Decompose n - 1 into (2^s) * d
	s = 0
	d = n - 1
	while d % 2 == 0:
		s += 1
		d //= 2

	# Determine the bases to use
	nm2 = n - 2
	if isinstance(checks, int):
		# Randomly choose 'checks' bases in range [2, n-2]
		bases = {random.randint(2, nm2) for _ in range(checks)}
	else:
		# Filter provided bases to be within range
		bases = {a for a in checks if a <= nm2}

	# Test n against all chosen bases
	for a in bases:
		if not CheckStrongProbablePrimeToBaseA2(a, d, s, n):
			return False

	return True


def TestCheckAgaistSimple():
	"""
	Compares Miller-Rabin results to simple trial division.
	Prints discrepancies or every 1000th number.
	"""
	n = 3
	# Deterministic base set from Sorenson & Webster (2015) assuming the extended Riemann hypothesis
	# works for all n <= 2^81
	limitedSet = {2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41}
	limitedSetLim = 2 ** 81

	while True:
		pp = MillerRabinPrimalityTest(n, limitedSet)
		np = SimplePrime(n)
		deterministic = n <= limitedSetLim
		if (n - 1) % 1000 == 0 or pp != np:
			print(n, pp, np, "deterministic" if deterministic else "probabilistic")
		if deterministic:
			assert pp == np
		n += 2  # Skip even numbers


def TestBiggerNums():
	"""
	Tests random odd integers up to 2^24 with Miller-Rabin.
	Reports if test is deterministic (n <= 2^81).
	"""
	limitedSet = {2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41}
	limitedSetLim = 2 ** 81

	while True:
		n = random.randint(3, 2 ** 24)
		if n % 2 == 0:
			continue
		deterministic = n <= limitedSetLim

		print("n", n)
		pp = MillerRabinPrimalityTest(n, limitedSet)
		print("is prime", pp, ",", "deterministic" if deterministic else "probabilistic")


def IsPowerOfTwo(n):
	while n > 2:
		if n % 2 != 0:
			return False
		n = n // 2
	return n % 2 == 0


def TestFermatPrimes():
	"""
	Tests numbers of the form 2^i + 1 (Fermat numbers) for primality.
	Reports whether they are actual Fermat primes.
	"""
	limitedSet = {2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41}
	limitedSetLim = 2 ** 81

	i = 2
	while True:
		n = 2 ** i + 1
		chk = MillerRabinPrimalityTest(n, limitedSet)
		deterministic = n <= limitedSetLim
		if chk:
			# Fermat primes occur only when i is a power of 2
			isFermatPrime = IsPowerOfTwo(i)
			print(
				i, n, chk,
				"is fermat prime" if isFermatPrime else "",
				"deterministic" if deterministic else "probabilistic"
			)
		i += 1


def LucasLehmerTest(p):
	s = 4
	M = 2**p - 1
	for i in range(p-2):
		s = ((s ** 2) - 2) % M
	return s == 0


def TestMersennePrimes():
	"""
	Tests numbers of the form 2^i - 1 (Mersenne numbers) for primality.
	"""

	i = 1
	while 1:
		n = 2 ** i - 1
		#chk = MillerRabinPrimalityTest(n) # Not efficient in this context
		chk = LucasLehmerTest(i)
		
		if chk:
			print (i, "2^{}-1".format(i), chk)
		i+=1


if __name__ == "__main__":
	TestCheckAgaistSimple()