bpo-46055: Speed up binary shifting operators

Lib/test/test_long.py

@@ -946,8 +946,13 @@ def test_huge_lshift(self, size):
         self.assertEqual(1 << (sys.maxsize + 1000), 1 << 1000 << sys.maxsize)
 
     def test_huge_rshift(self):
-        self.assertEqual(42 >> (1 << 1000), 0)
-        self.assertEqual((-42) >> (1 << 1000), -1)
+        huge_shift = 1 << 1000
+        self.assertEqual(42 >> huge_shift, 0)
+        self.assertEqual((-42) >> huge_shift, -1)
+        self.assertEqual(1123 >> huge_shift, 0)
+        self.assertEqual((-1123) >> huge_shift, -1)
+        self.assertEqual(2**128 >> huge_shift, 0)
+        self.assertEqual(-2**128 >> huge_shift, -1)
 
     @support.cpython_only
     @support.bigmemtest(sys.maxsize + 500, memuse=2/15, dry_run=False)
@@ -956,6 +961,60 @@ def test_huge_rshift_of_huge(self, size):
         self.assertEqual(huge >> (sys.maxsize + 1), (1 << 499) + 5)
         self.assertEqual(huge >> (sys.maxsize + 1000), 0)
 
+    def test_small_rshift(self):
+        self.assertEqual(42 >> 1, 21)
+        self.assertEqual((-42) >> 1, -21)
+        self.assertEqual(43 >> 1, 21)
+        self.assertEqual((-43) >> 1, -22)
+
+        self.assertEqual(1122 >> 1, 561)
+        self.assertEqual((-1122) >> 1, -561)
+        self.assertEqual(1123 >> 1, 561)
+        self.assertEqual((-1123) >> 1, -562)
+
+        self.assertEqual(2**128 >> 1, 2**127)
+        self.assertEqual(-2**128 >> 1, -2**127)
+        self.assertEqual((2**128 + 1) >> 1, 2**127)
+        self.assertEqual(-(2**128 + 1) >> 1, -2**127 - 1)
+
+    def test_medium_rshift(self):
+        self.assertEqual(42 >> 9, 0)
+        self.assertEqual((-42) >> 9, -1)
+        self.assertEqual(1122 >> 9, 2)
+        self.assertEqual((-1122) >> 9, -3)
+        self.assertEqual(2**128 >> 9, 2**119)
+        self.assertEqual(-2**128 >> 9, -2**119)
+
+    def test_big_rshift(self):
+        self.assertEqual(42 >> 32, 0)
+        self.assertEqual((-42) >> 32, -1)
+        self.assertEqual(1122 >> 32, 0)
+        self.assertEqual((-1122) >> 32, -1)
+        self.assertEqual(2**128 >> 32, 2**96)
+        self.assertEqual(-2**128 >> 32, -2**96)
+
+    def test_small_lshift(self):
+        self.assertEqual(42 << 1, 84)
+        self.assertEqual((-42) << 1, -84)
+        self.assertEqual(561 << 1, 1122)
+        self.assertEqual((-561) << 1, -1122)
+        self.assertEqual(2**127 << 1, 2**128)
+        self.assertEqual(-2**127 << 1, -2**128)
+
+    def test_medium_lshift(self):
+        self.assertEqual(42 << 9, 21504)
+        self.assertEqual((-42) << 9, -21504)
+        self.assertEqual(1122 << 9, 574464)
+        self.assertEqual((-1122) << 9, -574464)
+
+    def test_big_lshift(self):
+        self.assertEqual(42 << 32, 42 * 2**32)
+        self.assertEqual((-42) << 32, -42 * 2**32)
+        self.assertEqual(1122 << 32, 1122 * 2**32)
+        self.assertEqual((-1122) << 32, -1122 * 2**32)
+        self.assertEqual(2**128 << 32, 2**160)
+        self.assertEqual(-2**128 << 32, -2**160)
+
     @support.cpython_only
     def test_small_ints_in_huge_calculation(self):
         a = 2 ** 100

Objects/longobject.c

@@ -32,7 +32,9 @@ static inline stwodigits
 medium_value(PyLongObject *x)
 {
     assert(IS_MEDIUM_VALUE(x));
-    return ((stwodigits)Py_SIZE(x)) * x->ob_digit[0];
+    Py_ssize_t s = Py_SIZE(x);
+    stwodigits d = x->ob_digit[0];
+    return s == 1 ? d : s == -1 ? -d : 0;
 }
 
 #define IS_SMALL_INT(ival) (-_PY_NSMALLNEGINTS <= (ival) && (ival) < _PY_NSMALLPOSINTS)
@@ -4493,17 +4495,26 @@ long_rshift1(PyLongObject *a, Py_ssize_t wordshift, digit remshift)
     Py_ssize_t newsize, hishift, i, j;
     digit lomask, himask;
 
+    if (IS_MEDIUM_VALUE(a)) {
+        stwodigits m, x;
+        digit shift;
+        m = medium_value(a);
+        shift = wordshift == 0 ? remshift : PyLong_SHIFT;
+        x = m < 0 ? ~(~m >> shift) : m >> shift;
+        return _PyLong_FromSTwoDigits(x);
+    }
+
     if (Py_SIZE(a) < 0) {
         /* Right shifting negative numbers is harder */
         PyLongObject *a1, *a2;
-        a1 = (PyLongObject *) long_invert(a);
+        a1 = (PyLongObject *)long_invert(a);
         if (a1 == NULL)
             return NULL;
-        a2 = (PyLongObject *) long_rshift1(a1, wordshift, remshift);
+        a2 = (PyLongObject *)long_rshift1(a1, wordshift, remshift);
         Py_DECREF(a1);
         if (a2 == NULL)
             return NULL;
-        z = (PyLongObject *) long_invert(a2);
+        z = (PyLongObject *)long_invert(a2);
         Py_DECREF(a2);
     }
     else {
@@ -4518,8 +4529,8 @@ long_rshift1(PyLongObject *a, Py_ssize_t wordshift, digit remshift)
             return NULL;
         for (i = 0, j = wordshift; i < newsize; i++, j++) {
             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;
+            if (i + 1 < newsize)
+                z->ob_digit[i] |= (a->ob_digit[j + 1] << hishift) & himask;
         }
         z = maybe_small_long(long_normalize(z));
     }
@@ -4565,11 +4576,17 @@ _PyLong_Rshift(PyObject *a, size_t shiftby)
 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;
 
+    if (wordshift == 0 && IS_MEDIUM_VALUE(a)) {
+        stwodigits m = medium_value(a);
+        // bypass undefined shift operator behavior
+        stwodigits x = m < 0 ? -(-m << remshift) : m << remshift;
+        return _PyLong_FromSTwoDigits(x);
+    }
+
     oldsize = Py_ABS(Py_SIZE(a));
     newsize = oldsize + wordshift;
     if (remshift)