gh-90213: Speed up right shifts of negative integers

Lib/test/test_long.py

@@ -984,6 +984,10 @@ def test_medium_rshift(self):
         self.assertEqual((-1122) >> 9, -3)
         self.assertEqual(2**128 >> 9, 2**119)
         self.assertEqual(-2**128 >> 9, -2**119)
+        # Exercise corner case of the current algorithm, where the result of
+        # shifting a two-limb int by the limb size still has two limbs.
+        self.assertEqual((1 - BASE*BASE) >> SHIFT, -BASE)
+        self.assertEqual((BASE - 1 - BASE*BASE) >> SHIFT, -BASE)
 
     def test_big_rshift(self):
         self.assertEqual(42 >> 32, 0)

Misc/NEWS.d/next/Core and Builtins/2021-12-24-20-21-45.bpo-46055.R0QMVQ.rst

@@ -1,2 +1,3 @@
-Speed up shifting operation involving integers less than
-:c:macro:`PyLong_BASE`. Patch by Xinhang Xu.
+Speed up integer shifts: faster left and right shifts of integers smaller than
+:c:macro:`PyLong_BASE`, and faster right shifts in the general case. Patches by
+Xinhang Xu and Mark Dickinson.

Objects/longobject.c

@@ -4486,13 +4486,20 @@ divmod_shift(PyObject *shiftby, Py_ssize_t *wordshift, digit *remshift)
     return 0;
 }
 
+/* Inner function for both long_rshift and _PyLong_Rshift, shifting an
+   integer right by a strictly positive shift. */
+
 static PyObject *
 long_rshift1(PyLongObject *a, Py_ssize_t wordshift, digit remshift)
 {
     PyLongObject *z = NULL;
-    Py_ssize_t newsize, hishift, i, j;
+    Py_ssize_t newsize, hishift, i, j, size_a;
     twodigits accum;
+    int a_negative;
+
+    assert (wordshift > 0 || remshift > 0);
 
+    /* Fast path for small a. */
     if (IS_MEDIUM_VALUE(a)) {
         stwodigits m, x;
         digit shift;
@@ -4502,37 +4509,66 @@ long_rshift1(PyLongObject *a, Py_ssize_t wordshift, digit remshift)
         return _PyLong_FromSTwoDigits(x);
     }
 
-    if (Py_SIZE(a) < 0) {
-        /* Right shifting negative numbers is harder */
-        PyLongObject *a1, *a2;
-        a1 = (PyLongObject *) long_invert(a);
-        if (a1 == NULL)
-            return NULL;
-        a2 = (PyLongObject *) long_rshift1(a1, wordshift, remshift);
-        Py_DECREF(a1);
-        if (a2 == NULL)
-            return NULL;
-        z = (PyLongObject *) long_invert(a2);
-        Py_DECREF(a2);
+    a_negative = Py_SIZE(a) < 0;
+    size_a = Py_ABS(Py_SIZE(a));
+
+    if (a_negative) {
+        /* For shifting negative integers, it's convenient to adjust so that
+           0 < remshift <= PyLong_SHIFT. */
+        if (remshift == 0) {
+            remshift = PyLong_SHIFT;
+            wordshift -= 1;
+        }
+        assert(wordshift >= 0);
+    }
+
+    newsize = size_a - wordshift;
+    if (newsize <= 0) {
+        return PyLong_FromLong(-a_negative);
+    }
+    z = _PyLong_New(newsize);
+    if (z == NULL) {
+        return NULL;
+    }
+    hishift = PyLong_SHIFT - remshift;
+
+    if (a_negative) {
+        /*
+            For a positive integer a and nonnegative shift, we have:
+
+                (-a) >> shift == -((a + 2**shift - 1) >> shift).
+
+            In the addition `a + (2**shift - 1)`, the low `wordshift` digits of
+            `2**shift - 1` all have value `PyLong_MASK`, so we get a carry out
+            from the bottom `wordshift` digits when at least one of the least
+            significant `wordshift` digits of `a` is nonzero. Digit `wordshift`
+            of `2**shift - 1` has value `PyLong_MASK >> hishift`.
+        */
+        Py_SET_SIZE(z, -newsize);
+
+        digit sticky = 0;
+        for (j = 0; j < wordshift; j++) {
+            sticky |= a->ob_digit[j];
+        }
+        accum = a->ob_digit[j++] + (PyLong_MASK >> hishift)
+              + (digit)(sticky != 0);
     }
     else {
-        newsize = Py_SIZE(a) - wordshift;
-        if (newsize <= 0)
-            return PyLong_FromLong(0);
-        hishift = PyLong_SHIFT - remshift;
-        z = _PyLong_New(newsize);
-        if (z == NULL)
-            return NULL;
         j = wordshift;
-        accum = a->ob_digit[j++] >> remshift;
-        for (i = 0; j < Py_SIZE(a); i++, j++) {
-            accum |= (twodigits)a->ob_digit[j] << hishift;
-            z->ob_digit[i] = (digit)(accum & PyLong_MASK);
-            accum >>= PyLong_SHIFT;
-        }
-        z->ob_digit[i] = (digit)accum;
-        z = maybe_small_long(long_normalize(z));
+        accum = a->ob_digit[j++];
+    }
+
+    accum >>= remshift;
+    assert(j == wordshift+1);
+    for (i = 0; j < size_a; i++, j++) {
+        accum += (twodigits)a->ob_digit[j] << hishift;
+        z->ob_digit[i] = (digit)(accum & PyLong_MASK);
+        accum >>= PyLong_SHIFT;
     }
+    assert(i == newsize - 1);
+    assert(accum <= PyLong_MASK);
+    z->ob_digit[i] = (digit)accum;
+    z = maybe_small_long(long_normalize(z));
     return (PyObject *)z;
 }
 
@@ -4548,6 +4584,9 @@ long_rshift(PyObject *a, PyObject *b)
         PyErr_SetString(PyExc_ValueError, "negative shift count");
         return NULL;
     }
+    if (Py_SIZE(b) == 0) {
+        return long_long(a);
+    }
     if (Py_SIZE(a) == 0) {
         return PyLong_FromLong(0);
     }
@@ -4564,6 +4603,10 @@ _PyLong_Rshift(PyObject *a, size_t shiftby)
     digit remshift;
 
     assert(PyLong_Check(a));
+
+    if (shiftby == 0) {
+        return long_long(a);
+    }
     if (Py_SIZE(a) == 0) {
         return PyLong_FromLong(0);
     }