From 81e40f0446eaab13387ba90b747683f8819fd906 Mon Sep 17 00:00:00 2001 From: Serhiy Storchaka Date: Mon, 6 Jul 2026 20:29:09 +0300 Subject: [PATCH] gh-153222: Add math.integer.isprime() and math.integer.primes() Draft implementation of prime-number functionality for the math.integer module (deferred out of the initial PEP 791 scope), supporting integers less than 2**64. isprime() uses the deterministic Miller-Rabin test: bases {2, 7, 61} below 4759123141, Jim Sinclair's seven bases up to 2**64 (verified against the Feitsma-Galway list of base-2 strong pseudoprimes). The result is exact for the whole supported range; larger arguments raise OverflowError. primes(start=2, stop=None) returns an iterator of the primes in [start, stop), unbounded if stop is None. Neither function is mirrored into the math namespace (PEP 791). Co-Authored-By: Claude Fable 5 --- Doc/library/math.integer.rst | 37 +- Doc/whatsnew/3.16.rst | 9 + .../pycore_global_objects_fini_generated.h | 1 + Include/internal/pycore_global_strings.h | 1 + .../internal/pycore_runtime_init_generated.h | 1 + .../internal/pycore_unicodeobject_generated.h | 4 + Lib/test/test_math_integer.py | 290 ++++++++++++++ ...-07-06-14-00-00.gh-issue-153222.pRiMe1.rst | 3 + Modules/clinic/mathintegermodule.c.h | 111 +++++- Modules/mathintegermodule.c | 372 +++++++++++++++++- 10 files changed, 821 insertions(+), 8 deletions(-) create mode 100644 Misc/NEWS.d/next/Library/2026-07-06-14-00-00.gh-issue-153222.pRiMe1.rst diff --git a/Doc/library/math.integer.rst b/Doc/library/math.integer.rst index c3f34cdfd85410..d34e7f821c4df9 100644 --- a/Doc/library/math.integer.rst +++ b/Doc/library/math.integer.rst @@ -13,8 +13,9 @@ These functions accept integers and objects that implement the :meth:`~object.__index__` method which is used to convert the object to an integer number. -The following functions are provided by this module. All return values are -computed exactly and are integers. +The following functions are provided by this module. +Unless stated otherwise below, +all return values are computed exactly and are integers. .. function:: comb(n, k, /) @@ -46,6 +47,20 @@ computed exactly and are integers. returns ``0``. +.. function:: isprime(n, /) + + Return ``True`` if *n* is a prime number, ``False`` otherwise. + + A prime number is a natural number greater than 1 + that is not a product of two smaller natural numbers. + Negative numbers, ``0`` and ``1`` are not prime. + + The argument must be less than 2\ :sup:`64`; + raises :exc:`OverflowError` otherwise. + + .. versionadded:: next + + .. function:: isqrt(n, /) Return the integer square root of the nonnegative integer *n*. This is the @@ -83,3 +98,21 @@ computed exactly and are integers. and the function returns ``n!``. Raises :exc:`ValueError` if either of the arguments are negative. + + +.. function:: primes(start=2, stop=None) + + Return an iterator of the prime numbers *p* with ``start <= p < stop``, + in increasing order. + If *stop* is ``None`` (the default), the iteration does not stop. + + Roughly equivalent to + ``(p for p in itertools.count(start) if isprime(p))`` + for the unbounded form. + + The bounds must be less than 2\ :sup:`64`; + raises :exc:`OverflowError` otherwise. + Unbounded iteration raises :exc:`OverflowError` + if the candidates reach that limit. + + .. versionadded:: next diff --git a/Doc/whatsnew/3.16.rst b/Doc/whatsnew/3.16.rst index e8c75d2f571061..eb2537d15ce6e5 100644 --- a/Doc/whatsnew/3.16.rst +++ b/Doc/whatsnew/3.16.rst @@ -279,6 +279,15 @@ math (Contributed by Jeff Epler in :gh:`150534`.) +math.integer +------------ + +* Added :func:`math.integer.isprime` for primality testing of integers + less than 2**64 and + :func:`math.integer.primes` for iterating over prime numbers. + (Contributed by Serhiy Storchaka in :gh:`153222`.) + + os -- diff --git a/Include/internal/pycore_global_objects_fini_generated.h b/Include/internal/pycore_global_objects_fini_generated.h index 7d952a1e52561a..be16ea2d4aa522 100644 --- a/Include/internal/pycore_global_objects_fini_generated.h +++ b/Include/internal/pycore_global_objects_fini_generated.h @@ -2097,6 +2097,7 @@ _PyStaticObjects_CheckRefcnt(PyInterpreterState *interp) { _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(stdout)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(step)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(steps)); + _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(stop)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(store_name)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(strategy)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(strftime)); diff --git a/Include/internal/pycore_global_strings.h b/Include/internal/pycore_global_strings.h index 8a8bbc3b6d05bf..91f2b1c979ee3c 100644 --- a/Include/internal/pycore_global_strings.h +++ b/Include/internal/pycore_global_strings.h @@ -820,6 +820,7 @@ struct _Py_global_strings { STRUCT_FOR_ID(stdout) STRUCT_FOR_ID(step) STRUCT_FOR_ID(steps) + STRUCT_FOR_ID(stop) STRUCT_FOR_ID(store_name) STRUCT_FOR_ID(strategy) STRUCT_FOR_ID(strftime) diff --git a/Include/internal/pycore_runtime_init_generated.h b/Include/internal/pycore_runtime_init_generated.h index 366d2d300fb478..273164bc882688 100644 --- a/Include/internal/pycore_runtime_init_generated.h +++ b/Include/internal/pycore_runtime_init_generated.h @@ -2095,6 +2095,7 @@ extern "C" { INIT_ID(stdout), \ INIT_ID(step), \ INIT_ID(steps), \ + INIT_ID(stop), \ INIT_ID(store_name), \ INIT_ID(strategy), \ INIT_ID(strftime), \ diff --git a/Include/internal/pycore_unicodeobject_generated.h b/Include/internal/pycore_unicodeobject_generated.h index 00d6297432b5fc..179a47deb2dc89 100644 --- a/Include/internal/pycore_unicodeobject_generated.h +++ b/Include/internal/pycore_unicodeobject_generated.h @@ -3060,6 +3060,10 @@ _PyUnicode_InitStaticStrings(PyInterpreterState *interp) { _PyUnicode_InternStatic(interp, &string); assert(_PyUnicode_CheckConsistency(string, 1)); assert(PyUnicode_GET_LENGTH(string) != 1); + string = &_Py_ID(stop); + _PyUnicode_InternStatic(interp, &string); + assert(_PyUnicode_CheckConsistency(string, 1)); + assert(PyUnicode_GET_LENGTH(string) != 1); string = &_Py_ID(store_name); _PyUnicode_InternStatic(interp, &string); assert(_PyUnicode_CheckConsistency(string, 1)); diff --git a/Lib/test/test_math_integer.py b/Lib/test/test_math_integer.py index d1bc776dc42391..ef0ed0b7791a4f 100644 --- a/Lib/test/test_math_integer.py +++ b/Lib/test/test_math_integer.py @@ -1,5 +1,7 @@ from decimal import Decimal from fractions import Fraction +import itertools +import random import unittest from test import support @@ -78,6 +80,47 @@ def py_factorial(n): outer *= inner return outer << (n - count_set_bits(n)) +# Reference implementations for primality testing. + +def primes_below(n): + """List of the primes below n, by the sieve of Eratosthenes.""" + sieve = bytearray([1]) * n + sieve[:2] = bytes(2) + i = 2 + while i * i < n: + if sieve[i]: + sieve[i*i::i] = bytes(len(range(i*i, n, i))) + i += 1 + return [i for i in range(n) if sieve[i]] + +# The Miller-Rabin test with the first 13 primes as bases is known to be +# exact for n < 3.3 * 10**24. More bases are used for larger inputs, for +# which the test is probabilistic (but independent of the Baillie-PSW test +# used in the implementation). +MILLER_RABIN_BASES = primes_below(100) + +def py_isprime(n): + """Miller-Rabin primality test, for cross-checking.""" + if n < 2: + return False + for p in MILLER_RABIN_BASES: + if n % p == 0: + return n == p + d = n - 1 + s = (d & -d).bit_length() - 1 + d >>= s + for a in MILLER_RABIN_BASES: + x = pow(a, d, n) + if x == 1 or x == n - 1: + continue + for _ in range(s - 1): + x = x * x % n + if x == n - 1: + break + else: + return False + return True + class IntMathTests(unittest.TestCase): import math.integer as module @@ -406,6 +449,246 @@ class MathTests(IntMathTests): import math as module +# isprime() and primes() exist only in math.integer, not in math, so their +# tests are not in IntMathTests (which is re-run against math above). + +class IsPrimeTests(unittest.TestCase): + import math.integer as module + + def test_isprime_small(self): + isprime = self.module.isprime + sieve = set(primes_below(10**4)) + for n in range(-10, 10**4): + with self.subTest(n=n): + self.assertIs(isprime(n), n in sieve) + + def test_isprime_negative(self): + isprime = self.module.isprime + self.assertIs(isprime(-1), False) + self.assertIs(isprime(-2), False) + self.assertIs(isprime(-3), False) + self.assertIs(isprime(-10**100), False) + + def test_isprime_carmichael(self): + # Carmichael numbers are composite (A002997). + isprime = self.module.isprime + for n in [561, 1105, 1729, 2465, 2821, 6601, 8911, 10585, 15841, + 29341, 41041, 46657, 52633, 62745, 63973, 75361]: + with self.subTest(n=n): + self.assertIs(isprime(n), False) + + def test_isprime_strong_pseudoprimes_base_2(self): + # Composites that pass the strong probable prime test to base 2 + # (A001262); they must be caught by the strong Lucas test. + isprime = self.module.isprime + for n in [2047, 3277, 4033, 4681, 8321, 15841, 29341, 42799, 49141, + 52633, 65281, 74665, 80581, 85489, 88357, 90751, + 3825123056546413051]: + with self.subTest(n=n): + self.assertIs(isprime(n), False) + + def test_isprime_strong_lucas_pseudoprimes(self): + # Composites that pass the strong Lucas test with Selfridge's + # parameters (A217255). + isprime = self.module.isprime + for n in [5459, 5777, 10877, 16109, 18971, 22499, 24569, 25199, + 40309, 58519, 75077, 97439]: + with self.subTest(n=n): + self.assertIs(isprime(n), False) + + def test_isprime_base_divisors(self): + # Divisors of the Miller-Rabin bases exercise the case where a + # base is divisible by the tested number. + isprime = self.module.isprime + for base in [2, 7, 61, 325, 9375, 28178, 450775, 9780504, + 1795265022]: + d = 1 + while d * d <= base: + if base % d == 0: + for n in [d, base // d]: + with self.subTest(base=base, n=n): + self.assertIs(isprime(n), py_isprime(n)) + d += 1 + + def test_isprime_perfect_squares(self): + isprime = self.module.isprime + for k in range(1000): + with self.subTest(k=k): + self.assertIs(isprime(k*k), False) + for k in [2**31 - 1, 2**32 - 5]: + with self.subTest(k=k): + self.assertIs(isprime(k*k), False) + + def test_isprime_word_boundaries(self): + # Exercise the boundaries between the base sets. + isprime = self.module.isprime + for boundary in [2**32, 4759123141]: + for n in range(boundary - 200, boundary + 200): + with self.subTest(n=n): + self.assertIs(isprime(n), py_isprime(n)) + for n in range(2**64 - 200, 2**64): + with self.subTest(n=n): + self.assertIs(isprime(n), py_isprime(n)) + self.assertIs(isprime(2**64 - 59), True) # largest prime < 2**64 + + def test_isprime_large_values(self): + isprime = self.module.isprime + self.assertIs(isprime(2**61 - 1), True) # Mersenne prime + self.assertIs(isprime(2**62 - 1), False) + self.assertIs(isprime(10**19), False) + # Arguments not less than 2**64 are not supported. + for n in [2**64, 2**64 + 13, 2**89 - 1, 10**100]: + with self.subTest(n=n): + self.assertRaises(OverflowError, isprime, n) + + def test_isprime_random(self): + isprime = self.module.isprime + rng = random.Random(1729) + for bits in [32, 34, 63, 64]: + for _ in range(300): + n = rng.getrandbits(bits) + with self.subTest(n=n): + self.assertIs(isprime(n), py_isprime(n)) + for bits in [65, 80, 128]: + n = rng.getrandbits(bits) | (1 << (bits - 1)) + with self.subTest(n=n): + self.assertRaises(OverflowError, isprime, n) + + def test_isprime_integer_like(self): + isprime = self.module.isprime + self.assertIs(isprime(False), False) + self.assertIs(isprime(True), False) + self.assertIs(isprime(IntSubclass(7)), True) + self.assertIs(isprime(IntSubclass(8)), False) + self.assertIs(isprime(MyIndexable(97)), True) + self.assertIs(isprime(MyIndexable(-97)), False) + + def test_isprime_int_subclass_operators(self): + # Overridden operators of an int subclass must not affect the + # result. + isprime = self.module.isprime + self.assertIs(isprime(BadIntSubclass(97)), True) + self.assertIs(isprime(BadIntSubclass(2**61 - 1)), True) + self.assertIs(isprime(BadIntSubclass(2**62 - 1)), False) + + def test_isprime_non_integers(self): + isprime = self.module.isprime + for value in [7.0, 7.5, Decimal('7'), Fraction(7, 1), '7', 7.5j]: + with self.subTest(value=value): + self.assertRaises(TypeError, isprime, value) + self.assertRaises(TypeError, isprime) + self.assertRaises(TypeError, isprime, 7, 11) + + +class PrimesIterTests(unittest.TestCase): + import math.integer as module + + def test_primes(self): + primes = self.module.primes + expected = primes_below(10**4) + self.assertEqual(list(primes(stop=10**4)), expected) + self.assertEqual(len(expected), 1229) + self.assertEqual(list(itertools.islice(primes(), 25)), expected[:25]) + + def test_primes_start(self): + primes = self.module.primes + for start in [-10**100, -100, -1, 0, 1, 2]: + with self.subTest(start=start): + self.assertEqual(list(primes(start, 10)), [2, 3, 5, 7]) + self.assertEqual(list(primes(3, 10)), [3, 5, 7]) + self.assertEqual(list(primes(4, 10)), [5, 7]) + self.assertEqual(list(primes(8, 12)), [11]) + self.assertEqual(list(primes(9, 12)), [11]) + self.assertEqual(list(primes(7, 8)), [7]) + + def test_primes_stop(self): + primes = self.module.primes + # The range is half-open. + self.assertEqual(list(primes(2, 2)), []) + self.assertEqual(list(primes(2, 3)), [2]) + self.assertEqual(list(primes(3, 3)), []) + self.assertEqual(list(primes(7, 7)), []) + self.assertEqual(list(primes(2, -10)), []) + self.assertEqual(list(primes(10, 5)), []) + self.assertEqual(list(primes(stop=0)), []) + + def test_primes_unbounded(self): + primes = self.module.primes + it = primes() + self.assertEqual([next(it) for _ in range(5)], [2, 3, 5, 7, 11]) + it = primes(10**6) + self.assertEqual(next(it), 1000003) + + def test_primes_huge(self): + primes = self.module.primes + boundary = 10**18 + expected = [n for n in range(boundary - 200, boundary + 200) + if py_isprime(n)] + self.assertEqual(list(primes(boundary - 200, boundary + 200)), + expected) + # The top of the supported range. + expected = [n for n in range(2**64 - 200, 2**64) if py_isprime(n)] + self.assertEqual(list(primes(2**64 - 200, 2**64 - 1)), expected) + + def test_primes_overflow(self): + primes = self.module.primes + # The bounds must be less than 2**64. + self.assertRaises(OverflowError, primes, 2**64) + self.assertRaises(OverflowError, primes, 2**64 + 100, 2**64 + 200) + self.assertRaises(OverflowError, primes, 0, 2**64) + self.assertRaises(OverflowError, primes, 10**100) + # An unbounded iterator raises when it runs out of the + # supported range. + it = primes(2**64 - 60) + self.assertEqual(next(it), 2**64 - 59) + self.assertRaises(OverflowError, next, it) + self.assertRaises(StopIteration, next, it) + + def test_primes_iterator_protocol(self): + primes = self.module.primes + it = primes(2, 10) + self.assertIs(iter(it), it) + self.assertEqual(list(it), [2, 3, 5, 7]) + # An exhausted iterator stays exhausted. + self.assertEqual(list(it), []) + self.assertRaises(StopIteration, next, it) + # The iterator type cannot be instantiated directly. + self.assertRaises(TypeError, type(primes())) + + def test_primes_keywords(self): + primes = self.module.primes + self.assertEqual(list(primes(start=10, stop=30)), [11, 13, 17, 19, 23, 29]) + self.assertEqual(list(primes(10, stop=30)), [11, 13, 17, 19, 23, 29]) + + def test_primes_integer_like(self): + primes = self.module.primes + self.assertEqual(list(primes(True, 10)), [2, 3, 5, 7]) + self.assertEqual(list(primes(IntSubclass(3), IntSubclass(10))), [3, 5, 7]) + self.assertEqual(list(primes(MyIndexable(3), MyIndexable(10))), [3, 5, 7]) + # The yielded values are exact ints. + for p in primes(IntSubclass(3), IntSubclass(10)): + self.assertIs(type(p), int) + + def test_primes_int_subclass_operators(self): + # Overridden operators of an int subclass must not affect the + # iteration. + primes = self.module.primes + self.assertEqual(list(primes(BadIntSubclass(3), BadIntSubclass(10))), + [3, 5, 7]) + big = 10**18 + self.assertEqual(list(primes(BadIntSubclass(big), big + 100)), + [n for n in range(big, big + 100) if py_isprime(n)]) + + def test_primes_non_integers(self): + primes = self.module.primes + self.assertRaises(TypeError, primes, 2.5) + self.assertRaises(TypeError, primes, 2.5, 10) + self.assertRaises(TypeError, primes, 2, 10.5) + self.assertRaises(TypeError, primes, '2') + self.assertRaises(TypeError, primes, 2, '10') + self.assertRaises(TypeError, primes, 2, 10, 3) + + class MiscTests(unittest.TestCase): def test_module_name(self): @@ -416,6 +699,13 @@ def test_module_name(self): obj = getattr(math.integer, name) self.assertEqual(obj.__module__, 'math.integer') + def test_math_namespace(self): + # New functions are added only to math.integer, not to math + # (PEP 791). + import math + self.assertFalse(hasattr(math, 'isprime')) + self.assertFalse(hasattr(math, 'primes')) + if __name__ == '__main__': unittest.main() diff --git a/Misc/NEWS.d/next/Library/2026-07-06-14-00-00.gh-issue-153222.pRiMe1.rst b/Misc/NEWS.d/next/Library/2026-07-06-14-00-00.gh-issue-153222.pRiMe1.rst new file mode 100644 index 00000000000000..85891ca4150003 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2026-07-06-14-00-00.gh-issue-153222.pRiMe1.rst @@ -0,0 +1,3 @@ +Add :func:`math.integer.isprime` for primality testing of integers +less than 2**64 and +:func:`math.integer.primes` for iterating over prime numbers. diff --git a/Modules/clinic/mathintegermodule.c.h b/Modules/clinic/mathintegermodule.c.h index 29c2a0ac902ea3..e7796fc6c20baf 100644 --- a/Modules/clinic/mathintegermodule.c.h +++ b/Modules/clinic/mathintegermodule.c.h @@ -2,7 +2,11 @@ preserve [clinic start generated code]*/ -#include "pycore_modsupport.h" // _PyArg_CheckPositional() +#if defined(Py_BUILD_CORE) && !defined(Py_BUILD_CORE_MODULE) +# include "pycore_gc.h" // PyGC_Head +# include "pycore_runtime.h" // _Py_ID() +#endif +#include "pycore_modsupport.h" // _PyArg_UnpackKeywords() PyDoc_STRVAR(math_integer_gcd__doc__, "gcd($module, /, *integers)\n" @@ -67,6 +71,109 @@ PyDoc_STRVAR(math_integer_isqrt__doc__, #define MATH_INTEGER_ISQRT_METHODDEF \ {"isqrt", (PyCFunction)math_integer_isqrt, METH_O, math_integer_isqrt__doc__}, +PyDoc_STRVAR(math_integer_isprime__doc__, +"isprime($module, n, /)\n" +"--\n" +"\n" +"Return True if n is a prime number, False otherwise.\n" +"\n" +"The argument must be less than 2**64."); + +#define MATH_INTEGER_ISPRIME_METHODDEF \ + {"isprime", (PyCFunction)math_integer_isprime, METH_O, math_integer_isprime__doc__}, + +static int +math_integer_isprime_impl(PyObject *module, PyObject *n); + +static PyObject * +math_integer_isprime(PyObject *module, PyObject *n) +{ + PyObject *return_value = NULL; + int _return_value; + + _return_value = math_integer_isprime_impl(module, n); + if ((_return_value == -1) && PyErr_Occurred()) { + goto exit; + } + return_value = PyBool_FromLong((long)_return_value); + +exit: + return return_value; +} + +PyDoc_STRVAR(math_integer_primes__doc__, +"primes($module, /, start=2, stop=None)\n" +"--\n" +"\n" +"Return an iterator of the prime numbers in the range [start, stop).\n" +"\n" +"If stop is None, the iteration does not stop.\n" +"The bounds must be less than 2**64."); + +#define MATH_INTEGER_PRIMES_METHODDEF \ + {"primes", _PyCFunction_CAST(math_integer_primes), METH_FASTCALL|METH_KEYWORDS, math_integer_primes__doc__}, + +static PyObject * +math_integer_primes_impl(PyObject *module, PyObject *start, PyObject *stop); + +static PyObject * +math_integer_primes(PyObject *module, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames) +{ + PyObject *return_value = NULL; + #if defined(Py_BUILD_CORE) && !defined(Py_BUILD_CORE_MODULE) + + #define NUM_KEYWORDS 2 + static struct { + PyGC_Head _this_is_not_used; + PyObject_VAR_HEAD + Py_hash_t ob_hash; + PyObject *ob_item[NUM_KEYWORDS]; + } _kwtuple = { + .ob_base = PyVarObject_HEAD_INIT(&PyTuple_Type, NUM_KEYWORDS) + .ob_hash = -1, + .ob_item = { &_Py_ID(start), &_Py_ID(stop), }, + }; + #undef NUM_KEYWORDS + #define KWTUPLE (&_kwtuple.ob_base.ob_base) + + #else // !Py_BUILD_CORE + # define KWTUPLE NULL + #endif // !Py_BUILD_CORE + + static const char * const _keywords[] = {"start", "stop", NULL}; + static _PyArg_Parser _parser = { + .keywords = _keywords, + .fname = "primes", + .kwtuple = KWTUPLE, + }; + #undef KWTUPLE + PyObject *argsbuf[2]; + Py_ssize_t noptargs = nargs + (kwnames ? PyTuple_GET_SIZE(kwnames) : 0) - 0; + PyObject *start = NULL; + PyObject *stop = Py_None; + + args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, + /*minpos*/ 0, /*maxpos*/ 2, /*minkw*/ 0, /*varpos*/ 0, argsbuf); + if (!args) { + goto exit; + } + if (!noptargs) { + goto skip_optional_pos; + } + if (args[0]) { + start = args[0]; + if (!--noptargs) { + goto skip_optional_pos; + } + } + stop = args[1]; +skip_optional_pos: + return_value = math_integer_primes_impl(module, start, stop); + +exit: + return return_value; +} + PyDoc_STRVAR(math_integer_factorial__doc__, "factorial($module, n, /)\n" "--\n" @@ -156,4 +263,4 @@ math_integer_comb(PyObject *module, PyObject *const *args, Py_ssize_t nargs) exit: return return_value; } -/*[clinic end generated code: output=34697570c923a3af input=a9049054013a1b77]*/ +/*[clinic end generated code: output=172b302a40542e1c input=a9049054013a1b77]*/ diff --git a/Modules/mathintegermodule.c b/Modules/mathintegermodule.c index 0f660d461e349f..7978e8b5e9d494 100644 --- a/Modules/mathintegermodule.c +++ b/Modules/mathintegermodule.c @@ -5,12 +5,25 @@ #endif #include "Python.h" -#include "pycore_abstract.h" // _PyNumber_Index() -#include "pycore_bitutils.h" // _Py_bit_length() -#include "pycore_long.h" // _PyLong_GetZero() +#include "pycore_abstract.h" // _PyNumber_Index() +#include "pycore_bitutils.h" // _Py_bit_length() +#include "pycore_critical_section.h" // Py_BEGIN_CRITICAL_SECTION() +#include "pycore_long.h" // _PyLong_GetZero() #include "clinic/mathintegermodule.c.h" +typedef struct { + PyTypeObject *primes_iter_type; +} math_integer_state; + +static inline math_integer_state * +get_math_integer_state(PyObject *module) +{ + void *state = PyModule_GetState(module); + assert(state != NULL); + return (math_integer_state *)state; +} + /*[clinic input] module math module math.integer @@ -482,6 +495,323 @@ math_integer_isqrt(PyObject *module, PyObject *n) } +/* Primality testing with the deterministic Miller-Rabin test. + + A strong probable prime test is exact for n < 4759123141 with the + bases {2, 7, 61} (G. Jaeschke, "On strong pseudoprimes to several + bases", Mathematics of Computation 61 (1993), 915-926), and for all + n < 2**64 with the seven bases below, found by Jim Sinclair and + verified against the Feitsma-Galway exhaustive list of base-2 strong + pseudoprimes below 2**64 (see http://miller-rabin.appspot.com/ and + http://ntheory.org/pseudoprimes.html). +*/ + +static const uint8_t small_primes[] = { + 2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, +}; + +static const uint32_t mr_bases_small[] = {2, 7, 61}; +static const uint32_t mr_bases[] = { + 2, 325, 9375, 28178, 450775, 9780504, 1795265022, +}; + +#if defined(HAVE_GCC_UINT128_T) +static inline uint64_t +mulmod_u64(uint64_t a, uint64_t b, uint64_t n) +{ + return (uint64_t)((__uint128_t)a * b % n); +} +#elif defined(_MSC_VER) && defined(_M_X64) +# include +# pragma intrinsic(_umul128, _udiv128) +static inline uint64_t +mulmod_u64(uint64_t a, uint64_t b, uint64_t n) +{ + /* Since a, b < n, the high 64 bits of a*b are less than n and + _udiv128() cannot fault. */ + uint64_t hi, rem; + uint64_t lo = _umul128(a, b, &hi); + (void)_udiv128(hi, lo, n, &rem); + return rem; +} +#else +/* No fast 128-bit multiplication: fall back to shift-and-add, with the + modular additions written to avoid overflow for n close to 2**64. + Assumes a < n. */ +static inline uint64_t +mulmod_u64(uint64_t a, uint64_t b, uint64_t n) +{ + uint64_t result = 0; + while (b) { + if (b & 1) { + result = (result >= n - a) ? result - (n - a) : result + a; + } + a = (a >= n - a) ? a - (n - a) : a + a; + b >>= 1; + } + return result; +} +#endif + +static uint64_t +powmod_u64(uint64_t b, uint64_t e, uint64_t n) +{ + uint64_t result = 1; + b %= n; + while (e) { + if (e & 1) { + result = mulmod_u64(result, b, n); + } + b = mulmod_u64(b, b, n); + e >>= 1; + } + return result; +} + +/* Strong probable prime test to the given base. n must be odd and + coprime to the base (a base divisible by n gives no information and + is treated as passing). */ +static int +sprp_u64(uint64_t n, uint64_t base) +{ + base %= n; + if (base == 0) { + return 1; + } + uint64_t d = n - 1; + int s = 0; + while ((d & 1) == 0) { + d >>= 1; + s++; + } + uint64_t a = powmod_u64(base, d, n); + if (a == 1 || a == n - 1) { + return 1; + } + for (int r = 1; r < s; r++) { + a = mulmod_u64(a, a, n); + if (a == n - 1) { + return 1; + } + } + return 0; +} + +static int +isprime_u64(uint64_t n) +{ + for (size_t i = 0; i < Py_ARRAY_LENGTH(small_primes); i++) { + if (n % small_primes[i] == 0) { + return n == small_primes[i]; + } + } + if (n < 67 * 67) { + return n > 1; + } + const uint32_t *bases = mr_bases; + size_t num_bases = Py_ARRAY_LENGTH(mr_bases); + if (n < 4759123141) { + bases = mr_bases_small; + num_bases = Py_ARRAY_LENGTH(mr_bases_small); + } + for (size_t i = 0; i < num_bases; i++) { + if (!sprp_u64(n, bases[i])) { + return 0; + } + } + return 1; +} + +/*[clinic input] +math.integer.isprime -> bool + + n: object + / + +Return True if n is a prime number, False otherwise. + +The argument must be less than 2**64. +[clinic start generated code]*/ + +static int +math_integer_isprime_impl(PyObject *module, PyObject *n) +/*[clinic end generated code: output=c808dc14d8d86875 input=bab96f73b765d9cf]*/ +{ + n = _PyNumber_Index(n); + if (n == NULL) { + return -1; + } + int result; + uint64_t m; + if (_PyLong_IsNegative((PyLongObject *)n)) { + result = 0; + } + else if (PyLong_AsUInt64(n, &m) < 0) { + result = -1; + } + else { + result = isprime_u64(m); + } + Py_DECREF(n); + return result; +} + + +/* The iterator returned by primes(). */ + +typedef struct { + PyObject_HEAD + uint64_t candidate; /* next candidate to test; 2 or odd */ + uint64_t stop; /* exclusive upper bound; ignored if unbounded */ + char unbounded; + char done; + char overflowed; /* unbounded iteration ran out of the domain */ +} primesiterobject; + +static void +primes_iter_dealloc(PyObject *op) +{ + PyTypeObject *tp = Py_TYPE(op); + freefunc free_func = PyType_GetSlot(tp, Py_tp_free); + free_func(op); + Py_DECREF(tp); +} + +static PyObject * +primes_iter_next(PyObject *op) +{ + primesiterobject *it = (primesiterobject *)op; + PyObject *result = NULL; + Py_BEGIN_CRITICAL_SECTION(op); + for (;;) { + if (it->done) { + if (it->overflowed) { + /* The whole supported range has been iterated without + reaching a bound. Only reachable with stop=None. */ + it->overflowed = 0; + PyErr_SetString(PyExc_OverflowError, + "primes() argument must be less than 2**64"); + } + break; + } + if (PyErr_CheckSignals() < 0) { + break; + } + uint64_t cand = it->candidate; + if (!it->unbounded && cand >= it->stop) { + it->done = 1; + break; + } + /* Advance to the next candidate: 2 -> 3, odd -> odd + 2. */ + if (cand == 2) { + it->candidate = 3; + } + else if (cand == UINT64_MAX) { + it->done = 1; + it->overflowed = it->unbounded; + } + else { + it->candidate = cand + 2; + } + if (isprime_u64(cand)) { + result = PyLong_FromUnsignedLongLong(cand); + break; + } + } + Py_END_CRITICAL_SECTION(); + return result; +} + +static PyType_Slot primes_iter_slots[] = { + {Py_tp_dealloc, primes_iter_dealloc}, + {Py_tp_iter, PyObject_SelfIter}, + {Py_tp_iternext, primes_iter_next}, + {0, NULL}, +}; + +static PyType_Spec primes_iter_spec = { + .name = "math.integer.primes_iterator", + .basicsize = sizeof(primesiterobject), + .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_IMMUTABLETYPE + | Py_TPFLAGS_DISALLOW_INSTANTIATION), + .slots = primes_iter_slots, +}; + +/*[clinic input] +math.integer.primes + + start: object(c_default="NULL") = 2 + stop: object = None + +Return an iterator of the prime numbers in the range [start, stop). + +If stop is None, the iteration does not stop. +The bounds must be less than 2**64. +[clinic start generated code]*/ + +static PyObject * +math_integer_primes_impl(PyObject *module, PyObject *start, PyObject *stop) +/*[clinic end generated code: output=b92270dab45a6b84 input=b5abd91fb58a8d90]*/ +{ + uint64_t candidate = 2; + uint64_t stopval = 0; + char unbounded = 0; + char done = 0; + + if (stop == Py_None) { + unbounded = 1; + } + else { + PyObject *s = _PyNumber_Index(stop); + if (s == NULL) { + return NULL; + } + if (_PyLong_IsNegative((PyLongObject *)s)) { + done = 1; + } + else if (PyLong_AsUInt64(s, &stopval) < 0) { + Py_DECREF(s); + return NULL; + } + Py_DECREF(s); + } + if (start != NULL) { + PyObject *s = _PyNumber_Index(start); + if (s == NULL) { + return NULL; + } + if (_PyLong_IsNegative((PyLongObject *)s)) { + /* The first prime is 2. */ + } + else if (PyLong_AsUInt64(s, &candidate) < 0) { + Py_DECREF(s); + return NULL; + } + Py_DECREF(s); + if (candidate < 2) { + candidate = 2; + } + else if (candidate > 2) { + /* Round an even start up to the next odd number. */ + candidate |= 1; + } + } + + math_integer_state *state = get_math_integer_state(module); + primesiterobject *it = PyObject_New(primesiterobject, + state->primes_iter_type); + if (it == NULL) { + return NULL; + } + it->candidate = candidate; + it->stop = stopval; + it->unbounded = unbounded; + it->done = done; + it->overflowed = 0; + return (PyObject *)it; +} + + static unsigned long count_set_bits(unsigned long n) { @@ -1234,15 +1564,24 @@ static PyMethodDef math_integer_methods[] = { MATH_INTEGER_COMB_METHODDEF MATH_INTEGER_FACTORIAL_METHODDEF MATH_INTEGER_GCD_METHODDEF + MATH_INTEGER_ISPRIME_METHODDEF MATH_INTEGER_ISQRT_METHODDEF MATH_INTEGER_LCM_METHODDEF MATH_INTEGER_PERM_METHODDEF + MATH_INTEGER_PRIMES_METHODDEF {NULL, NULL} /* sentinel */ }; static int math_integer_exec(PyObject *module) { + math_integer_state *state = get_math_integer_state(module); + state->primes_iter_type = (PyTypeObject *)PyType_FromModuleAndSpec( + module, &primes_iter_spec, NULL); + if (state->primes_iter_type == NULL) { + return -1; + } + /* Fix the __name__ attribute of the module and the __module__ attribute * of its functions. */ @@ -1271,6 +1610,28 @@ math_integer_exec(PyObject *module) return 0; } +static int +math_integer_traverse(PyObject *module, visitproc visit, void *arg) +{ + math_integer_state *state = get_math_integer_state(module); + Py_VISIT(state->primes_iter_type); + return 0; +} + +static int +math_integer_clear(PyObject *module) +{ + math_integer_state *state = get_math_integer_state(module); + Py_CLEAR(state->primes_iter_type); + return 0; +} + +static void +math_integer_free(void *module) +{ + (void)math_integer_clear((PyObject *)module); +} + static PyModuleDef_Slot math_integer_slots[] = { _Py_ABI_SLOT, {Py_mod_exec, math_integer_exec}, @@ -1286,9 +1647,12 @@ static struct PyModuleDef math_integer_module = { PyModuleDef_HEAD_INIT, .m_name = "math.integer", .m_doc = module_doc, - .m_size = 0, + .m_size = sizeof(math_integer_state), .m_methods = math_integer_methods, .m_slots = math_integer_slots, + .m_traverse = math_integer_traverse, + .m_clear = math_integer_clear, + .m_free = math_integer_free, }; PyMODINIT_FUNC