bpo-29816: Shift operation now has less opportunity to raise OverflowError. (#680)

ValueError always is raised rather than OverflowError for negative counts.
Shifting zero with non-negative count always returns zero.

3 files changed, 82 insertions, 24 deletions

diff --git a/Lib/test/test_long.py b/Lib/test/test_long.py
index fd15f04..cc48259 100644
--- a/Lib/test/test_long.py
+++ b/Lib/test/test_long.py
@@ -906,11 +906,24 @@ class LongTest(unittest.TestCase):
             self.check_truediv(-x, y)
             self.check_truediv(-x, -y)
 
+    def test_negative_shift_count(self):
+        with self.assertRaises(ValueError):
+            42 << -3
+        with self.assertRaises(ValueError):
+            42 << -(1 << 1000)
+        with self.assertRaises(ValueError):
+            42 >> -3
+        with self.assertRaises(ValueError):
+            42 >> -(1 << 1000)
+
     def test_lshift_of_zero(self):
         self.assertEqual(0 << 0, 0)
         self.assertEqual(0 << 10, 0)
         with self.assertRaises(ValueError):
             0 << -1
+        self.assertEqual(0 << (1 << 1000), 0)
+        with self.assertRaises(ValueError):
+            0 << -(1 << 1000)
 
     @support.cpython_only
     def test_huge_lshift_of_zero(self):
@@ -918,8 +931,23 @@ class LongTest(unittest.TestCase):
         # Other implementations may have a different boundary for overflow,
         # or not raise at all.
         self.assertEqual(0 << sys.maxsize, 0)
-        with self.assertRaises(OverflowError):
-            0 << (sys.maxsize + 1)
+        self.assertEqual(0 << (sys.maxsize + 1), 0)
+
+    @support.cpython_only
+    @support.bigmemtest(sys.maxsize + 1000, memuse=2/15 * 2, dry_run=False)
+    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)
+
+    @support.cpython_only
+    @support.bigmemtest(sys.maxsize + 500, memuse=2/15, dry_run=False)
+    def test_huge_rshift_of_huge(self, size):
+        huge = ((1 << 500) + 11) << sys.maxsize
+        self.assertEqual(huge >> (sys.maxsize + 1), (1 << 499) + 5)
+        self.assertEqual(huge >> (sys.maxsize + 1000), 0)
 
     def test_small_ints(self):
         for i in range(-5, 257):
diff --git a/Misc/NEWS b/Misc/NEWS
index b22322f..2bea5ec 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -10,6 +10,10 @@ What's New in Python 3.7.0 alpha 1?
 Core and Builtins
 -----------------
 
+- bpo-29816: Shift operation now has less opportunity to raise OverflowError.
+  ValueError always is raised rather than OverflowError for negative counts.
+  Shifting zero with non-negative count always returns zero.
+
 - bpo-24821: Fixed the slowing down to 25 times in the searching of some
   unlucky Unicode characters.
 
diff --git a/Objects/longobject.c b/Objects/longobject.c
index 459eed9..4862b76 100644
--- a/Objects/longobject.c
+++ b/Objects/longobject.c
@@ -4277,15 +4277,54 @@ long_bool(PyLongObject *v)
     return Py_SIZE(v) != 0;
 }
 
+/* wordshift, remshift = divmod(shiftby, PyLong_SHIFT) */
+static int
+divmod_shift(PyLongObject *shiftby, Py_ssize_t *wordshift, digit *remshift)
+{
+    assert(PyLong_Check((PyObject *)shiftby));
+    assert(Py_SIZE(shiftby) >= 0);
+    Py_ssize_t lshiftby = PyLong_AsSsize_t((PyObject *)shiftby);
+    if (lshiftby >= 0) {
+        *wordshift = lshiftby / PyLong_SHIFT;
+        *remshift = lshiftby % PyLong_SHIFT;
+        return 0;
+    }
+    /* PyLong_Check(shiftby) is true and Py_SIZE(shiftby) >= 0, so it must
+       be that PyLong_AsSsize_t raised an OverflowError. */
+    assert(PyErr_ExceptionMatches(PyExc_OverflowError));
+    PyErr_Clear();
+    PyLongObject *wordshift_obj = divrem1(shiftby, PyLong_SHIFT, remshift);
+    if (wordshift_obj == NULL) {
+        return -1;
+    }
+    *wordshift = PyLong_AsSsize_t((PyObject *)wordshift_obj);
+    Py_DECREF(wordshift_obj);
+    if (*wordshift >= 0 && *wordshift < PY_SSIZE_T_MAX / (Py_ssize_t)sizeof(digit)) {
+        return 0;
+    }
+    PyErr_Clear();
+    /* Clip the value.  With such large wordshift the right shift
+       returns 0 and the left shift raises an error in _PyLong_New(). */
+    *wordshift = PY_SSIZE_T_MAX / sizeof(digit);
+    *remshift = 0;
+    return 0;
+}
+
 static PyObject *
 long_rshift(PyLongObject *a, PyLongObject *b)
 {
     PyLongObject *z = NULL;
-    Py_ssize_t shiftby, newsize, wordshift, loshift, hishift, i, j;
-    digit lomask, himask;
+    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;
+    }
+
     if (Py_SIZE(a) < 0) {
         /* Right shifting negative numbers is harder */
         PyLongObject *a1, *a2;
@@ -4300,19 +4339,11 @@ long_rshift(PyLongObject *a, PyLongObject *b)
         Py_DECREF(a2);
     }
     else {
-        shiftby = PyLong_AsSsize_t((PyObject *)b);
-        if (shiftby == -1L && PyErr_Occurred())
-            return NULL;
-        if (shiftby < 0) {
-            PyErr_SetString(PyExc_ValueError,
-                            "negative shift count");
+        if (divmod_shift(b, &wordshift, &loshift) < 0)
             return NULL;
-        }
-        wordshift = shiftby / PyLong_SHIFT;
-        newsize = Py_ABS(Py_SIZE(a)) - wordshift;
+        newsize = Py_SIZE(a) - wordshift;
         if (newsize <= 0)
             return PyLong_FromLong(0);
-        loshift = shiftby % PyLong_SHIFT;
         hishift = PyLong_SHIFT - loshift;
         lomask = ((digit)1 << hishift) - 1;
         himask = PyLong_MASK ^ lomask;
@@ -4336,27 +4367,22 @@ long_lshift(PyObject *v, PyObject *w)
     PyLongObject *a = (PyLongObject*)v;
     PyLongObject *b = (PyLongObject*)w;
     PyLongObject *z = NULL;
-    Py_ssize_t shiftby, oldsize, newsize, wordshift, remshift, i, j;
+    Py_ssize_t oldsize, newsize, wordshift, i, j;
+    digit remshift;
     twodigits accum;
 
     CHECK_BINOP(a, b);
 
-    shiftby = PyLong_AsSsize_t((PyObject *)b);
-    if (shiftby == -1L && PyErr_Occurred())
-        return NULL;
-    if (shiftby < 0) {
+    if (Py_SIZE(b) < 0) {
         PyErr_SetString(PyExc_ValueError, "negative shift count");
         return NULL;
     }
-
     if (Py_SIZE(a) == 0) {
         return PyLong_FromLong(0);
     }
 
-    /* wordshift, remshift = divmod(shiftby, PyLong_SHIFT) */
-    wordshift = shiftby / PyLong_SHIFT;
-    remshift  = shiftby - wordshift * PyLong_SHIFT;
-
+    if (divmod_shift(b, &wordshift, &remshift) < 0)
+        return NULL;
     oldsize = Py_ABS(Py_SIZE(a));
     newsize = oldsize + wordshift;
     if (remshift)