From b3b5f4735656a10edc9481ce059ae5d1913c8917 Mon Sep 17 00:00:00 2001 From: Sergey B Kirpichev Date: Fri, 22 May 2026 11:32:30 +0300 Subject: [PATCH 1/3] gh-90345: Add math.integer.isqrt_rem() --- Doc/library/math.integer.rst | 6 ++ Doc/whatsnew/3.16.rst | 7 ++ Lib/test/test_math_integer.py | 39 +++++++++ ...6-05-22-11-32-21.gh-issue-90345.Z9QjnG.rst | 1 + Modules/clinic/mathintegermodule.c.h | 11 ++- Modules/mathintegermodule.c | 85 ++++++++++++++++--- 6 files changed, 135 insertions(+), 14 deletions(-) create mode 100644 Misc/NEWS.d/next/Library/2026-05-22-11-32-21.gh-issue-90345.Z9QjnG.rst diff --git a/Doc/library/math.integer.rst b/Doc/library/math.integer.rst index c3f34cdfd85410c..29640190e140802 100644 --- a/Doc/library/math.integer.rst +++ b/Doc/library/math.integer.rst @@ -62,6 +62,12 @@ computed exactly and are integers. :trim: +.. function:: isqrt_rem(n, /) + + Return a pair of values (s,t) such that s=isqrt(n) and t=n-s*s. + The remainder *t* will be zero, if *n* is a perfect square. + + .. function:: lcm(*integers) Return the least common multiple of the specified integer arguments. diff --git a/Doc/whatsnew/3.16.rst b/Doc/whatsnew/3.16.rst index 2e5342e4f020534..989082a54fd15a1 100644 --- a/Doc/whatsnew/3.16.rst +++ b/Doc/whatsnew/3.16.rst @@ -86,6 +86,13 @@ New modules Improved modules ================ +math.integer +------------ + +* Add :func:`math.integer.isqrt_rem` to compute integer square root with + a remainder. + (Contributed by Sergey B Kirpichev in :gh:`90345`.) + os -- diff --git a/Lib/test/test_math_integer.py b/Lib/test/test_math_integer.py index 09a98d93bd636c9..559260a38983271 100644 --- a/Lib/test/test_math_integer.py +++ b/Lib/test/test_math_integer.py @@ -2,6 +2,7 @@ from fractions import Fraction import unittest from test import support +from math.integer import isqrt_rem class IntSubclass(int): @@ -249,6 +250,44 @@ def test_isqrt_huge(self, size): self.assertEqual(w.bit_length(), size // 2 + 1) self.assertEqual(w.bit_count(), 1) + def test_isqrt_rem(self): + test_values = ( + list(range(1000)) + + list(range(10**6 - 1000, 10**6 + 1000)) + + [2**e + i for e in range(60, 200) for i in range(-40, 40)] + + [3**9999, 10**5001] + ) + for value in test_values: + with self.subTest(value=value): + root, rem = isqrt_rem(value) + self.assertIs(type(root), int) + self.assertLessEqual(root*root, value) + self.assertLess(value, (root+1)*(root+1)) + self.assertIs(type(rem), int) + self.assertEqual(rem, value - root*root) + + # Negative values + with self.assertRaises(ValueError): + isqrt_rem(-1) + + # Integer-like things + self.assertEqual(isqrt_rem(True), (1, 0)) + self.assertEqual(isqrt_rem(False), (0, 0)) + self.assertEqual(isqrt_rem(MyIndexable(1729)), (41, 48)) + + with self.assertRaises(ValueError): + isqrt_rem(MyIndexable(-3)) + + # Non-integer-like things + bad_values = [ + 3.5, "a string", Decimal("3.5"), 3.5j, + 100.0, -4.0, + ] + for value in bad_values: + with self.subTest(value=value): + with self.assertRaises(TypeError): + isqrt_rem(value) + def test_perm(self): perm = self.module.perm factorial = self.module.factorial diff --git a/Misc/NEWS.d/next/Library/2026-05-22-11-32-21.gh-issue-90345.Z9QjnG.rst b/Misc/NEWS.d/next/Library/2026-05-22-11-32-21.gh-issue-90345.Z9QjnG.rst new file mode 100644 index 000000000000000..5c7b292af056221 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2026-05-22-11-32-21.gh-issue-90345.Z9QjnG.rst @@ -0,0 +1 @@ +Add :func:`math.integer.isqrt_rem`. diff --git a/Modules/clinic/mathintegermodule.c.h b/Modules/clinic/mathintegermodule.c.h index 29c2a0ac902ea39..a7929351d81b060 100644 --- a/Modules/clinic/mathintegermodule.c.h +++ b/Modules/clinic/mathintegermodule.c.h @@ -67,6 +67,15 @@ 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_isqrt_rem__doc__, +"isqrt_rem($module, n, /)\n" +"--\n" +"\n" +"Return a pair of values (s,t) such that s=isqrt(n) and t=n-s*s."); + +#define MATH_INTEGER_ISQRT_REM_METHODDEF \ + {"isqrt_rem", (PyCFunction)math_integer_isqrt_rem, METH_O, math_integer_isqrt_rem__doc__}, + PyDoc_STRVAR(math_integer_factorial__doc__, "factorial($module, n, /)\n" "--\n" @@ -156,4 +165,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=ec462d7c95cb54a8 input=a9049054013a1b77]*/ diff --git a/Modules/mathintegermodule.c b/Modules/mathintegermodule.c index cfad4154b2d3611..4428a89d9f4d319 100644 --- a/Modules/mathintegermodule.c +++ b/Modules/mathintegermodule.c @@ -341,18 +341,8 @@ _approximate_isqrt(uint64_t n) return (u << 15) + (uint32_t)((n >> 17) / u); } -/*[clinic input] -math.integer.isqrt - - n: object - / - -Return the integer part of the square root of the input. -[clinic start generated code]*/ - static PyObject * -math_integer_isqrt(PyObject *module, PyObject *n) -/*[clinic end generated code: output=551031e41a0f5d9e input=921ddd9853133d8d]*/ +_isqrt_rem(PyObject *n, PyObject **rem) { int a_too_large, c_bit_length; int64_t c, d; @@ -373,6 +363,9 @@ math_integer_isqrt(PyObject *module, PyObject *n) } if (_PyLong_IsZero((PyLongObject *)n)) { Py_DECREF(n); + if (rem) { + *rem = PyLong_FromLong(0); + } return PyLong_FromLong(0); } @@ -392,7 +385,15 @@ math_integer_isqrt(PyObject *module, PyObject *n) return NULL; } u = _approximate_isqrt(m << 2*shift) >> shift; - u -= (uint64_t)u * u > m; + uint64_t sq = (uint64_t)u * u; + u -= sq > m; + if (rem) { + if (sq > m) { + sq -= 2*(uint64_t)u + 1; + } + m -= sq; + *rem = PyLong_FromUnsignedLongLong(m); + } return PyLong_FromUnsignedLong(u); } @@ -460,14 +461,26 @@ math_integer_isqrt(PyObject *module, PyObject *n) goto error; } a_too_large = PyObject_RichCompareBool(n, b, Py_LT); - Py_DECREF(b); if (a_too_large == -1) { + Py_DECREF(b); goto error; } if (a_too_large) { + if (rem) { + Py_SETREF(b, PyNumber_Add(b, _PyLong_GetOne())); + Py_SETREF(b, PyNumber_Subtract(b, a)); + Py_SETREF(b, PyNumber_Subtract(b, a)); + } Py_SETREF(a, PyNumber_Subtract(a, _PyLong_GetOne())); } + if (!rem) { + Py_DECREF(b); + } + else { + Py_SETREF(b, PyNumber_Subtract(n, b)); + *rem = b; + } Py_DECREF(n); return a; @@ -478,6 +491,51 @@ math_integer_isqrt(PyObject *module, PyObject *n) } +/*[clinic input] +math.integer.isqrt + + n: object + / + +Return the integer part of the square root of the input. +[clinic start generated code]*/ + +static PyObject * +math_integer_isqrt(PyObject *module, PyObject *n) +/*[clinic end generated code: output=551031e41a0f5d9e input=921ddd9853133d8d]*/ +{ + return _isqrt_rem(n, NULL); +} + +/*[clinic input] +math.integer.isqrt_rem + + n: object + / + +Return a pair of values (s,t) such that s=isqrt(n) and t=n-s*s. +[clinic start generated code]*/ + +static PyObject * +math_integer_isqrt_rem(PyObject *module, PyObject *n) +/*[clinic end generated code: output=b17d11479d08cdc4 input=7ed2dd870818d2bb]*/ +{ + PyObject *rem = NULL; + PyObject *root = _isqrt_rem(n, &rem); + + if (root && rem) { + PyObject *tup = PyTuple_Pack(2, root, rem); + + Py_DECREF(root); + Py_DECREF(rem); + return tup; + } + Py_XDECREF(root); + Py_XDECREF(rem); + return NULL; +} + + static unsigned long count_set_bits(unsigned long n) { @@ -1231,6 +1289,7 @@ static PyMethodDef math_integer_methods[] = { MATH_INTEGER_FACTORIAL_METHODDEF MATH_INTEGER_GCD_METHODDEF MATH_INTEGER_ISQRT_METHODDEF + MATH_INTEGER_ISQRT_REM_METHODDEF MATH_INTEGER_LCM_METHODDEF MATH_INTEGER_PERM_METHODDEF {NULL, NULL} /* sentinel */ From cc6aca0c15c821321e5a4d95581e3326cabc8cf5 Mon Sep 17 00:00:00 2001 From: Sergey B Kirpichev Date: Sat, 23 May 2026 04:50:50 +0300 Subject: [PATCH 2/3] Apply suggestion from @picnixz MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Bénédikt Tran <10796600+picnixz@users.noreply.github.com> --- Modules/mathintegermodule.c | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/Modules/mathintegermodule.c b/Modules/mathintegermodule.c index 4428a89d9f4d319..4fc6cf153aae354 100644 --- a/Modules/mathintegermodule.c +++ b/Modules/mathintegermodule.c @@ -522,16 +522,14 @@ math_integer_isqrt_rem(PyObject *module, PyObject *n) { PyObject *rem = NULL; PyObject *root = _isqrt_rem(n, &rem); + PyObject *res = NULL; if (root && rem) { - PyObject *tup = PyTuple_Pack(2, root, rem); - - Py_DECREF(root); - Py_DECREF(rem); - return tup; + res = PyTuple_Pack(2, root, rem); } Py_XDECREF(root); Py_XDECREF(rem); + return res; return NULL; } From de60f6b497883db7eedb46f09f0a04deff2c49b1 Mon Sep 17 00:00:00 2001 From: Sergey B Kirpichev Date: Sat, 23 May 2026 05:27:52 +0300 Subject: [PATCH 3/3] address review: handle errors --- Modules/mathintegermodule.c | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/Modules/mathintegermodule.c b/Modules/mathintegermodule.c index 4fc6cf153aae354..bf7a2ce77bdb5b8 100644 --- a/Modules/mathintegermodule.c +++ b/Modules/mathintegermodule.c @@ -468,9 +468,20 @@ _isqrt_rem(PyObject *n, PyObject **rem) if (a_too_large) { if (rem) { - Py_SETREF(b, PyNumber_Add(b, _PyLong_GetOne())); - Py_SETREF(b, PyNumber_Subtract(b, a)); - Py_SETREF(b, PyNumber_Subtract(b, a)); + PyObject *tmp = PyNumber_Add(b, _PyLong_GetOne()); + + if (tmp == NULL) { + Py_DECREF(b); + goto error; + } + Py_SETREF(b, tmp); + tmp = PyNumber_Add(a, a); + if (tmp == NULL) { + Py_DECREF(b); + goto error; + } + Py_SETREF(b, PyNumber_Subtract(b, tmp)); + Py_DECREF(tmp); } Py_SETREF(a, PyNumber_Subtract(a, _PyLong_GetOne())); }