diff --git a/Lib/test/test_opcache.py b/Lib/test/test_opcache.py index 9480bf8b87baf3..2e0732be6fca77 100644 --- a/Lib/test/test_opcache.py +++ b/Lib/test/test_opcache.py @@ -1507,11 +1507,34 @@ def binary_op_bitwise_extend(): a, b = 3, 9 a ^= b self.assertEqual(a, 10) + a, b = 10, 2 + a = a >> b + self.assertEqual(a, 2) + 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") self.assert_no_opcode(binary_op_bitwise_extend, "BINARY_OP") + # check that after specialization of >> we handle negative shifts + for idx in range(100): + a, b = 2, 1 + if idx == 99: + b = -1 + try: + z = a >> b + except ValueError: + assert b == -1 + self.assertEqual(z, 1) + @cpython_only @requires_specialization def test_load_super_attr(self): diff --git a/Misc/NEWS.d/next/Core_and_Builtins/2025-01-30-20-15-44.gh-issue-100239.7_XP-Y.rst b/Misc/NEWS.d/next/Core_and_Builtins/2025-01-30-20-15-44.gh-issue-100239.7_XP-Y.rst new file mode 100644 index 00000000000000..d049549cc2ca98 --- /dev/null +++ b/Misc/NEWS.d/next/Core_and_Builtins/2025-01-30-20-15-44.gh-issue-100239.7_XP-Y.rst @@ -0,0 +1 @@ +Specialize ``BINARY_OP`` for left shift and right shift operations on compact ints. diff --git a/Python/specialize.c b/Python/specialize.c index b50728d4a2f3f3..eb404e29fe5b0c 100644 --- a/Python/specialize.c +++ b/Python/specialize.c @@ -237,6 +237,10 @@ _PyCode_Quicken(_Py_CODEUNIT *instructions, Py_ssize_t size, int enable_counters #define SPEC_FAIL_BINARY_OP_SUBSCR_BYTES 48 #define SPEC_FAIL_BINARY_OP_SUBSCR_STRUCTTIME 49 #define SPEC_FAIL_BINARY_OP_SUBSCR_RANGE 50 +#define SPEC_FAIL_BINARY_OP_LSHIFT_INT 51 +#define SPEC_FAIL_BINARY_OP_LSHIFT_DIFFERENT_TYPES 52 +#define SPEC_FAIL_BINARY_OP_RSHIFT_INT 53 +#define SPEC_FAIL_BINARY_OP_RSHIFT_DIFFERENT_TYPES 54 /* Calls */ @@ -1951,6 +1955,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: @@ -1978,6 +1988,12 @@ binary_op_fail_kind(int oparg, PyObject *lhs, PyObject *rhs) return SPEC_FAIL_BINARY_OP_REMAINDER; case NB_RSHIFT: case NB_INPLACE_RSHIFT: + if (!Py_IS_TYPE(lhs, Py_TYPE(rhs))) { + return SPEC_FAIL_BINARY_OP_RSHIFT_DIFFERENT_TYPES; + } + if (PyLong_CheckExact(lhs)) { + return SPEC_FAIL_BINARY_OP_RSHIFT_INT; + } return SPEC_FAIL_BINARY_OP_RSHIFT; case NB_SUBTRACT: case NB_INPLACE_SUBTRACT: @@ -2124,6 +2140,13 @@ is_compactlong(PyObject *v) _PyLong_IsCompact((PyLongObject *)v); } +static inline int +is_compactnonnegativelong(PyObject *v) +{ + return PyLong_CheckExact(v) && + _PyLong_IsNonNegativeCompact((PyLongObject *)v); +} + /* sequence * int helpers: bypass PyNumber_Multiply dispatch overhead by calling sq_repeat directly with PyLong_AsSsize_t. */ @@ -2180,6 +2203,16 @@ compactlongs_guard(PyObject *lhs, PyObject *rhs) return (is_compactlong(lhs) && is_compactlong(rhs)); } +static int +shift_guard(PyObject *lhs, PyObject *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) \ static PyObject * \ (NAME)(PyObject *lhs, PyObject *rhs) \ @@ -2191,8 +2224,22 @@ compactlongs_guard(PyObject *lhs, PyObject *rhs) BITWISE_LONGS_ACTION(compactlongs_or, |) BITWISE_LONGS_ACTION(compactlongs_and, &) BITWISE_LONGS_ACTION(compactlongs_xor, ^) +BITWISE_LONGS_ACTION(compactlongs_rshift, >>) #undef BITWISE_LONGS_ACTION +static PyObject * +compactlongs_lshift(PyObject *lhs, PyObject *rhs) +{ + // Left-shifting a negative signed value is undefined behavior in C, so + // perform the shift in unsigned arithmetic and reinterpret. The guard + // limits rhs_val to [0, 16] and compact long magnitudes fit comfortably + // in stwodigits with that headroom, so no overflow occurs. + Py_ssize_t rhs_val = _PyLong_CompactValue((PyLongObject *)rhs); + Py_ssize_t lhs_val = _PyLong_CompactValue((PyLongObject *)lhs); + stwodigits result = (stwodigits)((unsigned long long)lhs_val << rhs_val); + return PyLong_FromLongLong(result); +} + /* float-long */ static inline int @@ -2273,6 +2320,13 @@ static _PyBinaryOpSpecializationDescr binaryop_extend_descrs[] = { {NB_INPLACE_AND, compactlongs_guard, compactlongs_and, &PyLong_Type, 1, NULL, NULL}, {NB_INPLACE_XOR, compactlongs_guard, compactlongs_xor, &PyLong_Type, 1, NULL, NULL}, + /* long-long shifts: guards also check rhs is non-negative and <= 16 to + avoid undefined behavior and overflow, so type alone is not sufficient. */ + {NB_LSHIFT, shift_guard, compactlongs_lshift, &PyLong_Type, 1, NULL, NULL}, + {NB_RSHIFT, shift_guard, compactlongs_rshift, &PyLong_Type, 1, NULL, NULL}, + {NB_INPLACE_LSHIFT, shift_guard, compactlongs_lshift, &PyLong_Type, 1, NULL, NULL}, + {NB_INPLACE_RSHIFT, shift_guard, compactlongs_rshift, &PyLong_Type, 1, NULL, NULL}, + /* float-long arithmetic: guards also check NaN and compactness. */ {NB_ADD, float_compactlong_guard, float_compactlong_add, &PyFloat_Type, 1, NULL, NULL}, {NB_SUBTRACT, float_compactlong_guard, float_compactlong_subtract, &PyFloat_Type, 1, NULL, NULL},