beam: Miller-Rabin + Pollard rho factorization, brute-force verified

wizard-work/beam/factor.py

@@ -0,0 +1,124 @@
+"""Integer primality (deterministic Miller-Rabin for n < 2^64) and factorization
+(Pollard's rho with Brent cycle detection + deterministic Miller-Rabin).
+"""
+import math
+import random
+
+
+_SMALL_PRIMES = (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37)
+
+
+def is_prime(n: int) -> bool:
+    if n < 2:
+        return False
+    for p in _SMALL_PRIMES:
+        if n == p:
+            return True
+        if n % p == 0:
+            return False
+    d = n - 1
+    s = 0
+    while d % 2 == 0:
+        d //= 2
+        s += 1
+    for a in (2, 3, 5, 7, 11, 13, 17):
+        if a % n == 0:
+            continue
+        x = pow(a, d, n)
+        if x == 1 or x == n - 1:
+            continue
+        composite = True
+        for _ in range(s - 1):
+            x = x * x % n
+            if x == n - 1:
+                composite = False
+                break
+        if composite:
+            return False
+    return True
+
+
+def _pollard_brent(n: int) -> int:
+    if n % 2 == 0:
+        return 2
+    if n % 3 == 0:
+        return 3
+    while True:
+        c = random.randrange(1, n)
+        f = lambda x: (x * x + c) % n
+        y, m, g, r, q = random.randrange(2, n), 1, 1, 1, 1
+        while g == 1:
+            x = y
+            for _ in range(r):
+                y = f(y)
+            k = 0
+            while k < r and g == 1:
+                ys = y
+                for _ in range(min(m, r - k)):
+                    y = f(y)
+                    diff = x - y
+                    if diff < 0:
+                        diff = -diff
+                    q = q * diff % n
+                g = math.gcd(q, n)
+                k += m
+            r *= 2
+            m *= 2
+        if g == n:
+            while True:
+                ys = f(ys)
+                diff = x - ys
+                if diff < 0:
+                    diff = -diff
+                g = math.gcd(diff, n)
+                if g != 1:
+                    break
+        if g != n:
+            return g
+
+
+def factor(n: int) -> list[int]:
+    if n < 0:
+        return factor(-n)
+    if n < 2:
+        return []
+    n = abs(n)
+    factors = []
+
+    def _factor(m):
+        if m == 1:
+            return
+        if is_prime(m):
+            factors.append(m)
+            return
+        d = _pollard_brent(m)
+        _factor(d)
+        _factor(m // d)
+
+    for p in _SMALL_PRIMES:
+        while n % p == 0:
+            factors.append(p)
+            n //= p
+    if n > 1:
+        _factor(n)
+    factors.sort()
+    return factors
+
+
+def totient(n: int) -> int:
+    if n < 1:
+        raise ValueError("n must be >= 1")
+    if n == 1:
+        return 1
+    fs = factor(n)
+    result = n
+    for p in set(fs):
+        result -= result // p
+    return result
+
+
+if __name__ == "__main__":
+    import sys
+    for arg in sys.argv[1:]:
+        n = int(arg)
+        print(n, is_prime(n), factor(n))

wizard-work/beam/test_factor.py

@@ -0,0 +1,85 @@
+"""Tests for factor.py"""
+import random
+from factor import is_prime, factor, totient
+
+
+def _trial_prime(n):
+    if n < 2:
+        return False
+    if n < 4:
+        return True
+    if n % 2 == 0:
+        return False
+    i = 3
+    while i * i <= n:
+        if n % i == 0:
+            return False
+        i += 2
+    return True
+
+
+def test_small_primes():
+    for n in range(2, 3000):
+        assert is_prime(n) == _trial_prime(n), n
+
+
+def test_carmichaels():
+    for c in [561, 1105, 1729, 2465, 2821, 6601, 8911, 10585, 15841, 29341, 41041, 46657, 52633, 62745, 63973]:
+        assert not is_prime(c), c
+
+
+def test_large_known_primes():
+    for p in [2**31 - 1, 2**61 - 1, 999999999999999989, 32416190071]:
+        assert is_prime(p), p
+
+
+def test_factor_small():
+    for n in range(2, 1000):
+        fs = factor(n)
+        prod = 1
+        for f in fs:
+            prod *= f
+            assert is_prime(f)
+        assert prod == n
+        assert fs == sorted(fs)
+
+
+def test_factor_random_large():
+    random.seed(42)
+    for _ in range(300):
+        def rnd_prime(lo, hi):
+            while True:
+                x = random.randrange(lo, hi)
+                if is_prime(x):
+                    return x
+        a = rnd_prime(1000, 100000)
+        b = rnd_prime(1000, 100000)
+        n = a * b
+        fs = factor(n)
+        assert sorted(fs) == sorted([a, b]), (n, fs)
+
+
+def test_factor_powers():
+    for base in [2, 3, 5, 7, 11]:
+        for e in [1, 2, 3, 5, 8, 12]:
+            n = base ** e
+            fs = factor(n)
+            assert fs == [base] * e
+
+
+def test_totient():
+    assert totient(1) == 1
+    assert totient(2) == 1
+    assert totient(10) == 4
+    assert totient(12) == 4
+    assert totient(30) == 8
+
+
+def test_edge():
+    assert factor(0) == []
+    assert factor(1) == []
+    assert factor(-12) == [2, 2, 3]
+    assert is_prime(0) is False
+    assert is_prime(1) is False
+    assert is_prime(2) is True
+    assert is_prime(999999999999999937 * 2) is False