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

Misc/ACKS

@@ -1964,6 +1964,7 @@ Doug Wyatt
 Xiang Zhang
 Robert Xiao
 Florent Xicluna
+Xinhang Xu
 Arnon Yaari
 Alakshendra Yadav
 Hirokazu Yamamoto

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

@@ -0,0 +1,2 @@
+Speed up shifting operation involving integers less than
+:c:macro:`PyLong_BASE`. Patch by Xinhang Xu.

Objects/longobject.c

@@ -4493,6 +4493,15 @@ long_rshift1(PyLongObject *a, Py_ssize_t wordshift, digit remshift)
  Py_ssize_t newsize, hishift, i, j;
  twodigits accum;
 
+ 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;
@@ -4566,11 +4575,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)