From a5119e7d75c9729fc36c059d05f3d7132e7f6bb4 Mon Sep 17 00:00:00 2001 From: Serhiy Storchaka Date: Sun, 19 May 2019 14:14:38 +0300 Subject: [PATCH] bpo-36957: Add _PyLong_Rshift() and _PyLong_Lshift(). (GH-13416) --- Include/longobject.h | 3 ++ Modules/mathmodule.c | 30 +++--------- Objects/floatobject.c | 4 +- Objects/longobject.c | 103 ++++++++++++++++++++++++++++++------------ 4 files changed, 87 insertions(+), 53 deletions(-) diff --git a/Include/longobject.h b/Include/longobject.h index b696f544b9c..a24bbea3a90 100644 --- a/Include/longobject.h +++ b/Include/longobject.h @@ -230,6 +230,9 @@ PyAPI_FUNC(PyObject *) _PyLong_GCD(PyObject *, PyObject *); #ifndef Py_LIMITED_API PyAPI_DATA(PyObject *) _PyLong_Zero; PyAPI_DATA(PyObject *) _PyLong_One; + +PyAPI_FUNC(PyObject *) _PyLong_Rshift(PyObject *, size_t); +PyAPI_FUNC(PyObject *) _PyLong_Lshift(PyObject *, size_t); #endif #ifdef __cplusplus diff --git a/Modules/mathmodule.c b/Modules/mathmodule.c index 821309221f8..7a0044a9fcf 100644 --- a/Modules/mathmodule.c +++ b/Modules/mathmodule.c @@ -1671,18 +1671,13 @@ math_isqrt(PyObject *module, PyObject *n) } d = 0; while (--s >= 0) { - PyObject *q, *shift; + PyObject *q; 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); + q = _PyLong_Rshift(n, 2U*c - d - e + 1U); if (q == NULL) { goto error; } @@ -1692,13 +1687,7 @@ math_isqrt(PyObject *module, PyObject *n) } /* 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); + Py_SETREF(a, _PyLong_Lshift(a, d - 1U - e)); if (a == NULL) { Py_DECREF(q); goto error; @@ -1939,9 +1928,9 @@ static PyObject * math_factorial(PyObject *module, PyObject *arg) /*[clinic end generated code: output=6686f26fae00e9ca input=6d1c8105c0d91fb4]*/ { - long x; + long x, two_valuation; int overflow; - PyObject *result, *odd_part, *two_valuation, *pyint_form; + PyObject *result, *odd_part, *pyint_form; if (PyFloat_Check(arg)) { PyObject *lx; @@ -1990,13 +1979,8 @@ math_factorial(PyObject *module, PyObject *arg) odd_part = factorial_odd_part(x); if (odd_part == NULL) return NULL; - two_valuation = PyLong_FromLong(x - count_set_bits(x)); - if (two_valuation == NULL) { - Py_DECREF(odd_part); - return NULL; - } - result = PyNumber_Lshift(odd_part, two_valuation); - Py_DECREF(two_valuation); + two_valuation = x - count_set_bits(x); + result = _PyLong_Lshift(odd_part, two_valuation); Py_DECREF(odd_part); return result; } diff --git a/Objects/floatobject.c b/Objects/floatobject.c index adb9b80c271..4ff43bb338f 100644 --- a/Objects/floatobject.c +++ b/Objects/floatobject.c @@ -463,13 +463,13 @@ float_richcompare(PyObject *v, PyObject *w, int op) */ PyObject *temp; - temp = PyNumber_Lshift(ww, _PyLong_One); + temp = _PyLong_Lshift(ww, 1); if (temp == NULL) goto Error; Py_DECREF(ww); ww = temp; - temp = PyNumber_Lshift(vv, _PyLong_One); + temp = _PyLong_Lshift(vv, 1); if (temp == NULL) goto Error; Py_DECREF(vv); diff --git a/Objects/longobject.c b/Objects/longobject.c index 9fb1fb02c27..1934328820c 100644 --- a/Objects/longobject.c +++ b/Objects/longobject.c @@ -4416,9 +4416,9 @@ long_bool(PyLongObject *v) /* wordshift, remshift = divmod(shiftby, PyLong_SHIFT) */ static int -divmod_shift(PyLongObject *shiftby, Py_ssize_t *wordshift, digit *remshift) +divmod_shift(PyObject *shiftby, Py_ssize_t *wordshift, digit *remshift) { - assert(PyLong_Check((PyObject *)shiftby)); + assert(PyLong_Check(shiftby)); assert(Py_SIZE(shiftby) >= 0); Py_ssize_t lshiftby = PyLong_AsSsize_t((PyObject *)shiftby); if (lshiftby >= 0) { @@ -4430,7 +4430,7 @@ divmod_shift(PyLongObject *shiftby, Py_ssize_t *wordshift, digit *remshift) be that PyLong_AsSsize_t raised an OverflowError. */ assert(PyErr_ExceptionMatches(PyExc_OverflowError)); PyErr_Clear(); - PyLongObject *wordshift_obj = divrem1(shiftby, PyLong_SHIFT, remshift); + PyLongObject *wordshift_obj = divrem1((PyLongObject *)shiftby, PyLong_SHIFT, remshift); if (wordshift_obj == NULL) { return -1; } @@ -4448,19 +4448,11 @@ divmod_shift(PyLongObject *shiftby, Py_ssize_t *wordshift, digit *remshift) } static PyObject * -long_rshift(PyLongObject *a, PyLongObject *b) +long_rshift1(PyLongObject *a, Py_ssize_t wordshift, digit remshift) { PyLongObject *z = NULL; - Py_ssize_t newsize, wordshift, hishift, i, j; - digit loshift, lomask, himask; - - CHECK_BINOP(a, b); - - if (Py_SIZE(b) < 0) { - PyErr_SetString(PyExc_ValueError, - "negative shift count"); - return NULL; - } + Py_ssize_t newsize, hishift, i, j; + digit lomask, himask; if (Py_SIZE(a) < 0) { /* Right shifting negative numbers is harder */ @@ -4468,7 +4460,7 @@ long_rshift(PyLongObject *a, PyLongObject *b) a1 = (PyLongObject *) long_invert(a); if (a1 == NULL) return NULL; - a2 = (PyLongObject *) long_rshift(a1, b); + a2 = (PyLongObject *) long_rshift1(a1, wordshift, remshift); Py_DECREF(a1); if (a2 == NULL) return NULL; @@ -4476,19 +4468,17 @@ long_rshift(PyLongObject *a, PyLongObject *b) Py_DECREF(a2); } else { - if (divmod_shift(b, &wordshift, &loshift) < 0) - return NULL; newsize = Py_SIZE(a) - wordshift; if (newsize <= 0) return PyLong_FromLong(0); - hishift = PyLong_SHIFT - loshift; + hishift = PyLong_SHIFT - remshift; lomask = ((digit)1 << hishift) - 1; himask = PyLong_MASK ^ lomask; z = _PyLong_New(newsize); if (z == NULL) return NULL; for (i = 0, j = wordshift; i < newsize; i++, j++) { - z->ob_digit[i] = (a->ob_digit[j] >> loshift) & lomask; + z->ob_digit[i] = (a->ob_digit[j] >> remshift) & lomask; if (i+1 < newsize) z->ob_digit[i] |= (a->ob_digit[j+1] << hishift) & himask; } @@ -4498,15 +4488,10 @@ long_rshift(PyLongObject *a, PyLongObject *b) } static PyObject * -long_lshift(PyObject *v, PyObject *w) +long_rshift(PyObject *a, PyObject *b) { - /* This version due to Tim Peters */ - PyLongObject *a = (PyLongObject*)v; - PyLongObject *b = (PyLongObject*)w; - PyLongObject *z = NULL; - Py_ssize_t oldsize, newsize, wordshift, i, j; + Py_ssize_t wordshift; digit remshift; - twodigits accum; CHECK_BINOP(a, b); @@ -4517,9 +4502,35 @@ long_lshift(PyObject *v, PyObject *w) if (Py_SIZE(a) == 0) { return PyLong_FromLong(0); } - if (divmod_shift(b, &wordshift, &remshift) < 0) return NULL; + return long_rshift1((PyLongObject *)a, wordshift, remshift); +} + +/* Return a >> shiftby. */ +PyObject * +_PyLong_Rshift(PyObject *a, size_t shiftby) +{ + Py_ssize_t wordshift; + digit remshift; + + assert(PyLong_Check(a)); + if (Py_SIZE(a) == 0) { + return PyLong_FromLong(0); + } + wordshift = shiftby / PyLong_SHIFT; + remshift = shiftby % PyLong_SHIFT; + return long_rshift1((PyLongObject *)a, wordshift, remshift); +} + +static PyObject * +long_lshift1(PyLongObject *a, Py_ssize_t wordshift, digit remshift) +{ + /* This version due to Tim Peters */ + PyLongObject *z = NULL; + Py_ssize_t oldsize, newsize, i, j; + twodigits accum; + oldsize = Py_ABS(Py_SIZE(a)); newsize = oldsize + wordshift; if (remshift) @@ -4547,6 +4558,42 @@ long_lshift(PyObject *v, PyObject *w) return (PyObject *) maybe_small_long(z); } +static PyObject * +long_lshift(PyObject *a, PyObject *b) +{ + Py_ssize_t wordshift; + digit remshift; + + CHECK_BINOP(a, b); + + if (Py_SIZE(b) < 0) { + PyErr_SetString(PyExc_ValueError, "negative shift count"); + return NULL; + } + if (Py_SIZE(a) == 0) { + return PyLong_FromLong(0); + } + if (divmod_shift(b, &wordshift, &remshift) < 0) + return NULL; + return long_lshift1((PyLongObject *)a, wordshift, remshift); +} + +/* Return a << shiftby. */ +PyObject * +_PyLong_Lshift(PyObject *a, size_t shiftby) +{ + Py_ssize_t wordshift; + digit remshift; + + assert(PyLong_Check(a)); + if (Py_SIZE(a) == 0) { + return PyLong_FromLong(0); + } + wordshift = shiftby / PyLong_SHIFT; + remshift = shiftby % PyLong_SHIFT; + return long_lshift1((PyLongObject *)a, wordshift, remshift); +} + /* Compute two's complement of digit vector a[0:m], writing result to z[0:m]. The digit vector a need not be normalized, but should not be entirely zero. a and z may point to the same digit vector. */ @@ -5552,7 +5599,7 @@ static PyNumberMethods long_as_number = { (inquiry)long_bool, /*tp_bool*/ (unaryfunc)long_invert, /*nb_invert*/ long_lshift, /*nb_lshift*/ - (binaryfunc)long_rshift, /*nb_rshift*/ + long_rshift, /*nb_rshift*/ long_and, /*nb_and*/ long_xor, /*nb_xor*/ long_or, /*nb_or*/