bpo-36957: Add _PyLong_Rshift() and _PyLong_Lshift(). (GH-13416)

This commit is contained in:
Serhiy Storchaka 2019-05-19 14:14:38 +03:00 committed by GitHub
parent 1d5bdef550
commit a5119e7d75
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 87 additions and 53 deletions

View File

@ -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

View File

@ -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;
}

View File

@ -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);

View File

@ -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*/