bpo-36887: add math.isqrt (GH-13244)
* Add math.isqrt function computing the integer square root.
* Code cleanup: remove redundant comments, rename some variables.
* Tighten up code a bit more; use Py_XDECREF to simplify error handling.
* Update Modules/mathmodule.c
Co-Authored-By: Serhiy Storchaka <storchaka@gmail.com>
* Update Modules/mathmodule.c
Use real argument clinic type instead of an alias
Co-Authored-By: Serhiy Storchaka <storchaka@gmail.com>
* Add proof sketch
* Updates from review.
* Correct and expand documentation.
* Fix bad reference handling on error; make some variables block-local; other tidying.
* Style and consistency fixes.
* Add missing error check; don't try to DECREF a NULL a
* Simplify some error returns.
* Another two test cases:
- clarify that floats are rejected even if they happen to be
squares of small integers
- TypeError beats ValueError for a negative float
* Documentation and markup improvements; thanks Serhiy for the suggestions!
* Cleaner Misc/NEWS entry wording.
* Clean up (with one fix) to the algorithm explanation and proof.
diff --git a/Modules/clinic/mathmodule.c.h b/Modules/clinic/mathmodule.c.h
index 1806a01..e677bd8 100644
--- a/Modules/clinic/mathmodule.c.h
+++ b/Modules/clinic/mathmodule.c.h
@@ -65,6 +65,15 @@
#define MATH_FSUM_METHODDEF \
{"fsum", (PyCFunction)math_fsum, METH_O, math_fsum__doc__},
+PyDoc_STRVAR(math_isqrt__doc__,
+"isqrt($module, n, /)\n"
+"--\n"
+"\n"
+"Return the integer part of the square root of the input.");
+
+#define MATH_ISQRT_METHODDEF \
+ {"isqrt", (PyCFunction)math_isqrt, METH_O, math_isqrt__doc__},
+
PyDoc_STRVAR(math_factorial__doc__,
"factorial($module, x, /)\n"
"--\n"
@@ -628,4 +637,4 @@
exit:
return return_value;
}
-/*[clinic end generated code: output=96e71135dce41c48 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=aeed62f403b90199 input=a9049054013a1b77]*/
diff --git a/Modules/mathmodule.c b/Modules/mathmodule.c
index 8f6a303..8213092 100644
--- a/Modules/mathmodule.c
+++ b/Modules/mathmodule.c
@@ -1476,6 +1476,266 @@
return count;
}
+/* Integer square root
+
+Given a nonnegative integer `n`, we want to compute the largest integer
+`a` for which `a * a <= n`, or equivalently the integer part of the exact
+square root of `n`.
+
+We use an adaptive-precision pure-integer version of Newton's iteration. Given
+a positive integer `n`, the algorithm produces at each iteration an integer
+approximation `a` to the square root of `n >> s` for some even integer `s`,
+with `s` decreasing as the iterations progress. On the final iteration, `s` is
+zero and we have an approximation to the square root of `n` itself.
+
+At every step, the approximation `a` is strictly within 1.0 of the true square
+root, so we have
+
+ (a - 1)**2 < (n >> s) < (a + 1)**2
+
+After the final iteration, a check-and-correct step is needed to determine
+whether `a` or `a - 1` gives the desired integer square root of `n`.
+
+The algorithm is remarkable in its simplicity. There's no need for a
+per-iteration check-and-correct step, and termination is straightforward: the
+number of iterations is known in advance (it's exactly `floor(log2(log2(n)))`
+for `n > 1`). The only tricky part of the correctness proof is in establishing
+that the bound `(a - 1)**2 < (n >> s) < (a + 1)**2` is maintained from one
+iteration to the next. A sketch of the proof of this is given below.
+
+In addition to the proof sketch, a formal, computer-verified proof
+of correctness (using Lean) of an equivalent recursive algorithm can be found
+here:
+
+ https://github.com/mdickinson/snippets/blob/master/proofs/isqrt/src/isqrt.lean
+
+
+Here's Python code equivalent to the C implementation below:
+
+ def isqrt(n):
+ """
+ Return the integer part of the square root of the input.
+ """
+ n = operator.index(n)
+
+ if n < 0:
+ raise ValueError("isqrt() argument must be nonnegative")
+ if n == 0:
+ return 0
+
+ c = (n.bit_length() - 1) // 2
+ a = 1
+ d = 0
+ for s in reversed(range(c.bit_length())):
+ e = d
+ d = c >> s
+ a = (a << d - e - 1) + (n >> 2*c - e - d + 1) // a
+ assert (a-1)**2 < n >> 2*(c - d) < (a+1)**2
+
+ return a - (a*a > n)
+
+
+Sketch of proof of correctness
+------------------------------
+
+The delicate part of the correctness proof is showing that the loop invariant
+is preserved from one iteration to the next. That is, just before the line
+
+ a = (a << d - e - 1) + (n >> 2*c - e - d + 1) // a
+
+is executed in the above code, we know that
+
+ (1) (a - 1)**2 < (n >> 2*(c - e)) < (a + 1)**2.
+
+(since `e` is always the value of `d` from the previous iteration). We must
+prove that after that line is executed, we have
+
+ (a - 1)**2 < (n >> 2*(c - d)) < (a + 1)**2
+
+To faciliate the proof, we make some changes of notation. Write `m` for
+`n >> 2*(c-d)`, and write `b` for the new value of `a`, so
+
+ b = (a << d - e - 1) + (n >> 2*c - e - d + 1) // a
+
+or equivalently:
+
+ (2) b = (a << d - e - 1) + (m >> d - e + 1) // a
+
+Then we can rewrite (1) as:
+
+ (3) (a - 1)**2 < (m >> 2*(d - e)) < (a + 1)**2
+
+and we must show that (b - 1)**2 < m < (b + 1)**2.
+
+From this point on, we switch to mathematical notation, so `/` means exact
+division rather than integer division and `^` is used for exponentiation. We
+use the `√` symbol for the exact square root. In (3), we can remove the
+implicit floor operation to give:
+
+ (4) (a - 1)^2 < m / 4^(d - e) < (a + 1)^2
+
+Taking square roots throughout (4), scaling by `2^(d-e)`, and rearranging gives
+
+ (5) 0 <= | 2^(d-e)a - √m | < 2^(d-e)
+
+Squaring and dividing through by `2^(d-e+1) a` gives
+
+ (6) 0 <= 2^(d-e-1) a + m / (2^(d-e+1) a) - √m < 2^(d-e-1) / a
+
+We'll show below that `2^(d-e-1) <= a`. Given that, we can replace the
+right-hand side of (6) with `1`, and now replacing the central
+term `m / (2^(d-e+1) a)` with its floor in (6) gives
+
+ (7) -1 < 2^(d-e-1) a + m // 2^(d-e+1) a - √m < 1
+
+Or equivalently, from (2):
+
+ (7) -1 < b - √m < 1
+
+and rearranging gives that `(b-1)^2 < m < (b+1)^2`, which is what we needed
+to prove.
+
+We're not quite done: we still have to prove the inequality `2^(d - e - 1) <=
+a` that was used to get line (7) above. From the definition of `c`, we have
+`4^c <= n`, which implies
+
+ (8) 4^d <= m
+
+also, since `e == d >> 1`, `d` is at most `2e + 1`, from which it follows
+that `2d - 2e - 1 <= d` and hence that
+
+ (9) 4^(2d - 2e - 1) <= m
+
+Dividing both sides by `4^(d - e)` gives
+
+ (10) 4^(d - e - 1) <= m / 4^(d - e)
+
+But we know from (4) that `m / 4^(d-e) < (a + 1)^2`, hence
+
+ (11) 4^(d - e - 1) < (a + 1)^2
+
+Now taking square roots of both sides and observing that both `2^(d-e-1)` and
+`a` are integers gives `2^(d - e - 1) <= a`, which is what we needed. This
+completes the proof sketch.
+
+*/
+
+/*[clinic input]
+math.isqrt
+
+ n: object
+ /
+
+Return the integer part of the square root of the input.
+[clinic start generated code]*/
+
+static PyObject *
+math_isqrt(PyObject *module, PyObject *n)
+/*[clinic end generated code: output=35a6f7f980beab26 input=5b6e7ae4fa6c43d6]*/
+{
+ int a_too_large, s;
+ size_t c, d;
+ PyObject *a = NULL, *b;
+
+ n = PyNumber_Index(n);
+ if (n == NULL) {
+ return NULL;
+ }
+
+ if (_PyLong_Sign(n) < 0) {
+ PyErr_SetString(
+ PyExc_ValueError,
+ "isqrt() argument must be nonnegative");
+ goto error;
+ }
+ if (_PyLong_Sign(n) == 0) {
+ Py_DECREF(n);
+ return PyLong_FromLong(0);
+ }
+
+ c = _PyLong_NumBits(n);
+ if (c == (size_t)(-1)) {
+ goto error;
+ }
+ c = (c - 1U) / 2U;
+
+ /* s = c.bit_length() */
+ s = 0;
+ while ((c >> s) > 0) {
+ ++s;
+ }
+
+ a = PyLong_FromLong(1);
+ if (a == NULL) {
+ goto error;
+ }
+ d = 0;
+ while (--s >= 0) {
+ PyObject *q, *shift;
+ size_t e = d;
+
+ d = c >> s;
+
+ /* q = (n >> 2*c - e - d + 1) // a */
+ shift = PyLong_FromSize_t(2U*c - d - e + 1U);
+ if (shift == NULL) {
+ goto error;
+ }
+ q = PyNumber_Rshift(n, shift);
+ Py_DECREF(shift);
+ if (q == NULL) {
+ goto error;
+ }
+ Py_SETREF(q, PyNumber_FloorDivide(q, a));
+ if (q == NULL) {
+ goto error;
+ }
+
+ /* a = (a << d - 1 - e) + q */
+ shift = PyLong_FromSize_t(d - 1U - e);
+ if (shift == NULL) {
+ Py_DECREF(q);
+ goto error;
+ }
+ Py_SETREF(a, PyNumber_Lshift(a, shift));
+ Py_DECREF(shift);
+ if (a == NULL) {
+ Py_DECREF(q);
+ goto error;
+ }
+ Py_SETREF(a, PyNumber_Add(a, q));
+ Py_DECREF(q);
+ if (a == NULL) {
+ goto error;
+ }
+ }
+
+ /* The correct result is either a or a - 1. Figure out which, and
+ decrement a if necessary. */
+
+ /* a_too_large = n < a * a */
+ b = PyNumber_Multiply(a, a);
+ if (b == NULL) {
+ goto error;
+ }
+ a_too_large = PyObject_RichCompareBool(n, b, Py_LT);
+ Py_DECREF(b);
+ if (a_too_large == -1) {
+ goto error;
+ }
+
+ if (a_too_large) {
+ Py_SETREF(a, PyNumber_Subtract(a, _PyLong_One));
+ }
+ Py_DECREF(n);
+ return a;
+
+ error:
+ Py_XDECREF(a);
+ Py_DECREF(n);
+ return NULL;
+}
+
/* Divide-and-conquer factorial algorithm
*
* Based on the formula and pseudo-code provided at:
@@ -2737,6 +2997,7 @@
MATH_ISFINITE_METHODDEF
MATH_ISINF_METHODDEF
MATH_ISNAN_METHODDEF
+ MATH_ISQRT_METHODDEF
MATH_LDEXP_METHODDEF
{"lgamma", math_lgamma, METH_O, math_lgamma_doc},
MATH_LOG_METHODDEF