handle case of negative shift

python · eendebakpt · Jan 29, 2025 · Jan 29, 2025 · Jan 29, 2025 · Jan 29, 2025
commit ca60ccc34046a166866fabc88c8f91b729cd0ef4
diff --git a/Lib/test/test_opcache.py b/Lib/test/test_opcache.py
@@ -1425,6 +1425,16 @@ def binary_op_bitwise_extend():
         self.assert_specialized(binary_op_bitwise_extend, "BINARY_OP_EXTEND")
         self.assert_no_opcode(binary_op_bitwise_extend, "BINARY_OP")
 
+        for idx in range(100):
+            a, b = 2, 1
+            if idx == 99:
+                b = -1
+            try:
+                a >> b
+            except ValueError:
+                assert b == -1
+            self.assertEqual(a, 1)
+
     @cpython_only
     @requires_specialization_ft
     def test_load_super_attr(self):

diff --git a/Python/specialize.c b/Python/specialize.c
@@ -2447,12 +2447,25 @@ is_compactlong(PyObject *v)
            _PyLong_IsCompact((PyLongObject *)v);
 }
 
+static inline int
+is_compactnonnegativelong(PyObject *v)
+{
+    return PyLong_CheckExact(v) &&
+           _PyLong_IsNonNegativeCompact((PyLongObject *)v);
+}
+
 static int
 compactlongs_guard(PyObject *lhs, PyObject *rhs)
 {
     return (is_compactlong(lhs) && is_compactlong(rhs));
 }
 
+static int
+compactlong_compactnonnegativelong_guard(PyObject *lhs, PyObject *rhs)
+{
+    return (is_compactlong(lhs) && is_compactnonnegativelong(rhs);
+}
+
 #define BITWISE_LONGS_ACTION(NAME, OP) \
     static PyObject * \
     (NAME)(PyObject *lhs, PyObject *rhs) \
@@ -2541,11 +2554,11 @@ static _PyBinaryOpSpecializationDescr compactlongs_specs[NB_OPARG_LAST+1] = {
     [NB_OR] = {compactlongs_guard, compactlongs_or},
     [NB_AND] = {compactlongs_guard, compactlongs_and},
     [NB_XOR] = {compactlongs_guard, compactlongs_xor},
-    [NB_RSHIFT] = {compactlongs_guard, compactlongs_rshift},
+    [NB_RSHIFT] = {compactlong_compactnonnegativelong_guard, compactlongs_rshift},
     [NB_INPLACE_OR] = {compactlongs_guard, compactlongs_or},
     [NB_INPLACE_AND] = {compactlongs_guard, compactlongs_and},
     [NB_INPLACE_XOR] = {compactlongs_guard, compactlongs_xor},
-    [NB_INPLACE_RSHIFT] = {compactlongs_guard, compactlongs_rshift},
+    [NB_INPLACE_RSHIFT] = {compactlong_compactnonnegativelong_guard, compactlongs_rshift},
 };
 
 static _PyBinaryOpSpecializationDescr float_compactlong_specs[NB_OPARG_LAST+1] = {