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

4 files changed, 87 insertions, 53 deletions

diff --git a/Include/longobject.h b/Include/longobject.h
index b696f54..a24bbea 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 8213092..7a0044a 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 adb9b80..4ff43bb 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 9fb1fb0..1934328 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*/