gh-100239: specialize bytes binary operations

e0173f14842aec44b42f9e1f32161b3cd4264ba4
commit e0173f14842aec44b42f9e1f32161b3cd4264ba4
diff --git a/Lib/test/test_opcache.py b/Lib/test/test_opcache.py
@@ -1439,6 +1439,21 @@ def binary_op_add_extend_sequences():
         self.assert_specialized(binary_op_add_extend_sequences, "BINARY_OP_EXTEND")
         self.assert_no_opcode(binary_op_add_extend_sequences, "BINARY_OP")
 
+        def binary_op_extend_bytes():
+            b1 = b"foo"
+            b2 = b"bar"
+            for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
+                bytes_sum = b1 + b2
+                self.assertEqual(bytes_sum, b"foobar")
+                bytes_repeat = b1 * 2
+                self.assertEqual(bytes_repeat, b"foofoo")
+                bytes_repeat = 2 * b1
+                self.assertEqual(bytes_repeat, b"foofoo")
+
+        binary_op_extend_bytes()
+        self.assert_specialized(binary_op_extend_bytes, "BINARY_OP_EXTEND")
+        self.assert_no_opcode(binary_op_extend_bytes, "BINARY_OP")
+
         def binary_op_zero_division():
             def compactlong_lhs(arg):
                 42 / arg
@@ -1953,6 +1968,29 @@ def binary_subscr_str_int_non_compact():
         self.assert_specialized(binary_subscr_str_int_non_compact, "BINARY_OP_SUBSCR_USTR_INT")
         self.assert_no_opcode(binary_subscr_str_int_non_compact, "BINARY_OP_SUBSCR_STR_INT")
 
+        def binary_subscr_bytes_int():
+            a = b"foobar"
+            for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
+                for idx, expected in enumerate(a):
+                    self.assertEqual(a[idx], expected)
+                self.assertEqual(a[-1], ord("r"))
+
+        binary_subscr_bytes_int()
+        self.assert_specialized(binary_subscr_bytes_int, "BINARY_OP_EXTEND")
+        self.assert_no_opcode(binary_subscr_bytes_int, "BINARY_OP")
+
+        def binary_subscr_bytes_slice():
+            a = b"foobar"
+            for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
+                self.assertEqual(a[1:4], b"oob")
+                self.assertEqual(a[::-1], b"raboof")
+                self.assertEqual(a[:], a)
+                self.assertEqual(a[10:20], b"")
+
+        binary_subscr_bytes_slice()
+        self.assert_specialized(binary_subscr_bytes_slice, "BINARY_OP_EXTEND")
+        self.assert_no_opcode(binary_subscr_bytes_slice, "BINARY_OP")
+
         def binary_subscr_getitems():
             class C:
                 def __init__(self, val):

diff --git a/Python/specialize.c b/Python/specialize.c
@@ -2193,6 +2193,70 @@ BITWISE_LONGS_ACTION(compactlongs_and, &)
 BITWISE_LONGS_ACTION(compactlongs_xor, ^)
 #undef BITWISE_LONGS_ACTION
 
+/* bytes subscripting */
+
+static inline int
+bytes_compactlong_guard(PyObject *lhs, PyObject *rhs)
+{
+    if (!PyBytes_CheckExact(lhs) || !is_compactlong(rhs)) {
+        return 0;
+    }
+    Py_ssize_t index = _PyLong_CompactValue((PyLongObject *)rhs);
+    Py_ssize_t size = PyBytes_GET_SIZE(lhs);
+    if (index < 0) {
+        index += size;
+    }
+    return 0 <= index && index < size;
+}
+
+static PyObject *
+bytes_compactlong_subscr(PyObject *lhs, PyObject *rhs)
+{
+    Py_ssize_t index = _PyLong_CompactValue((PyLongObject *)rhs);
+    Py_ssize_t size = PyBytes_GET_SIZE(lhs);
+    if (index < 0) {
+        index += size;
+    }
+    unsigned char value = (unsigned char)PyBytes_AS_STRING(lhs)[index];
+    return _PyLong_FromUnsignedChar(value);
+}
+
+static PyObject *
+bytes_slice_subscr(PyObject *lhs, PyObject *rhs)
+{
+    assert(PyBytes_CheckExact(lhs));
+    assert(PySlice_Check(rhs));
+
+    Py_ssize_t start, stop, step;
+    if (PySlice_Unpack(rhs, &start, &stop, &step) < 0) {
+        return NULL;
+    }
+    Py_ssize_t size = PyBytes_GET_SIZE(lhs);
+    Py_ssize_t slicelength = PySlice_AdjustIndices(size, &start, &stop, step);
+    if (slicelength <= 0) {
+        return Py_GetConstant(Py_CONSTANT_EMPTY_BYTES);
+    }
+    if (start == 0 && step == 1 && slicelength == size) {
+        return Py_NewRef(lhs);
+    }
+    if (step == 1) {
+        return PyBytes_FromStringAndSize(
+            PyBytes_AS_STRING(lhs) + start,
+            slicelength);
+    }
+
+    const char *source_buf = PyBytes_AS_STRING(lhs);
+    PyObject *result = PyBytes_FromStringAndSize(NULL, slicelength);
+    if (result == NULL) {
+        return NULL;
+    }
+    char *result_buf = PyBytes_AS_STRING(result);
+    for (Py_ssize_t cur = start, i = 0; i < slicelength; cur += step, i++) {
+        result_buf[i] = source_buf[cur];
+    }
+    return result;
+}
+
 /* float-long */
 
 static inline int
@@ -2273,6 +2337,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},
 
+    /* bytes[int]: guard includes bounds checks so the generic opcode still
+       raises IndexError for out-of-range indexes. */
+    {NB_SUBSCR, bytes_compactlong_guard, bytes_compactlong_subscr, &PyLong_Type, 1, NULL, NULL},
+    /* bytes[slice]: may return the original bytes object or the empty bytes
+       singleton, so the result is not guaranteed to be unique. */
+    {NB_SUBSCR, NULL, bytes_slice_subscr, &PyBytes_Type, 0, &PyBytes_Type, &PySlice_Type},
+
     /* 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},