add lshift; avoid undefined behavior

python · eendebakpt · Jan 29, 2025 · Jan 29, 2025 · Jan 29, 2025 · Jan 29, 2025
commit 004dce0408176668327ef2c8592249a33b0150d5
diff --git a/Lib/test/test_opcache.py b/Lib/test/test_opcache.py
@@ -1422,6 +1422,12 @@ def binary_op_bitwise_extend():
                 a, b = 10, 2
                 a >>= b
                 self.assertEqual(a, 2)
+                a, b = 10, 2
+                a = a << b
+                self.assertEqual(a, 40)
+                a, b = 10, 2
+                a <<= b
+                self.assertEqual(a, 40)
 
         binary_op_bitwise_extend()
         self.assert_specialized(binary_op_bitwise_extend, "BINARY_OP_EXTEND")

diff --git a/Python/specialize.c b/Python/specialize.c
@@ -585,6 +585,8 @@ _PyCode_Quicken(_Py_CODEUNIT *instructions, Py_ssize_t size, PyObject *consts,
 #define SPEC_FAIL_BINARY_OP_OR_DIFFERENT_TYPES          30
 #define SPEC_FAIL_BINARY_OP_XOR_INT                     31
 #define SPEC_FAIL_BINARY_OP_XOR_DIFFERENT_TYPES         32
+#define SPEC_FAIL_BINARY_OP_LSHIFT_INT                  33
+#define SPEC_FAIL_BINARY_OP_LSHIFT_DIFFERENT_TYPES      34
 #define SPEC_FAIL_BINARY_OP_RSHIFT_INT                  33
 #define SPEC_FAIL_BINARY_OP_RSHIFT_DIFFERENT_TYPES      34
 
@@ -2373,6 +2375,12 @@ binary_op_fail_kind(int oparg, PyObject *lhs, PyObject *rhs)
             return SPEC_FAIL_BINARY_OP_FLOOR_DIVIDE;
         case NB_LSHIFT:
         case NB_INPLACE_LSHIFT:
+            if (!Py_IS_TYPE(lhs, Py_TYPE(rhs))) {
+                return SPEC_FAIL_BINARY_OP_LSHIFT_DIFFERENT_TYPES;
+            }
+            if (PyLong_CheckExact(lhs)) {
+                return SPEC_FAIL_BINARY_OP_LSHIFT_INT;
+            }
             return SPEC_FAIL_BINARY_OP_LSHIFT;
         case NB_MATRIX_MULTIPLY:
         case NB_INPLACE_MATRIX_MULTIPLY:
@@ -2461,9 +2469,13 @@ compactlongs_guard(PyObject *lhs, PyObject *rhs)
 }
 
 static int
-compactlong_compactnonnegativelong_guard(PyObject *lhs, PyObject *rhs)
+shift_guard(PyObject *lhs, PyObject *rhs)
 {
-    return (is_compactlong(lhs) && is_compactnonnegativelong(rhs));
+    // we could use _long_is_small_int here, which is slightly faster than is_compactnonnegativelong
+
+    // rshift with value larger the the number of bits is undefined in C
+    // for lshift we do not want to overflow, but we always have at least 16 bits available
+    return (is_compactlong(lhs) && is_compactnonnegativelong(rhs) && (_PyLong_CompactValue((PyLongObject *)rhs) <= 16) );
 }
 
 #define BITWISE_LONGS_ACTION(NAME, OP) \
@@ -2477,6 +2489,7 @@ compactlong_compactnonnegativelong_guard(PyObject *lhs, PyObject *rhs)
 BITWISE_LONGS_ACTION(compactlongs_or, |)
 BITWISE_LONGS_ACTION(compactlongs_and, &)
 BITWISE_LONGS_ACTION(compactlongs_xor, ^)
+BITWISE_LONGS_ACTION(compactlongs_lshift, <<)
 BITWISE_LONGS_ACTION(compactlongs_rshift, >>)
 #undef BITWISE_LONGS_ACTION
 
@@ -2554,11 +2567,13 @@ 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] = {compactlong_compactnonnegativelong_guard, compactlongs_rshift},
+    [NB_LSHIFT] = {shift_guard, compactlongs_lshift},
+    [NB_RSHIFT] = {shift_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] = {compactlong_compactnonnegativelong_guard, compactlongs_rshift},
+    [NB_INPLACE_LSHIFT] = {shift_guard, compactlongs_lshift},
+    [NB_INPLACE_RSHIFT] = {shift_guard, compactlongs_rshift},
 };
 
 static _PyBinaryOpSpecializationDescr float_compactlong_specs[NB_OPARG_LAST+1] = {