gh-100239: specialize bytes binary operations

Include/internal/pycore_bytesobject.h

@@ -19,6 +19,13 @@ extern PyObject* _PyBytes_FormatEx(
  * reference rather than modifying its first argument in place. */
 extern PyObject* _PyBytes_Concat(PyObject *a, PyObject *b);
 
+extern int _PyBytes_OffsetFromIndex(
+    PyObject *op,
+    Py_ssize_t index,
+    Py_ssize_t *offset);
+
+extern PyObject* _PyBytes_SubscriptIndex(PyObject *op, Py_ssize_t index);
+
 extern PyObject* _PyBytes_FromHex(
     PyObject *string,
     int use_bytearray);

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,17 @@ 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_getitems():
             class C:
                 def __init__(self, val):

Objects/bytesobject.c

@@ -1728,6 +1728,34 @@ bytes_hash(PyObject *self)
     return hash;
 }
 
+int
+_PyBytes_OffsetFromIndex(PyObject *op, Py_ssize_t index, Py_ssize_t *offset)
+{
+    PyBytesObject *self = _PyBytes_CAST(op);
+    Py_ssize_t size = PyBytes_GET_SIZE(self);
+    if (index < 0) {
+        index += size;
+    }
+    if (index < 0 || index >= size) {
+        return 0;
+    }
+    *offset = index;
+    return 1;
+}
+
+PyObject *
+_PyBytes_SubscriptIndex(PyObject *op, Py_ssize_t index)
+{
+    PyBytesObject *self = _PyBytes_CAST(op);
+    Py_ssize_t offset;
+    if (!_PyBytes_OffsetFromIndex(op, index, &offset)) {
+        PyErr_SetString(PyExc_IndexError,
+                        "index out of range");
+        return NULL;
+    }
+    return _PyLong_FromUnsignedChar((unsigned char)self->ob_sval[offset]);
+}
+
 static PyObject*
 bytes_subscript(PyObject *op, PyObject* item)
 {
@@ -1736,14 +1764,7 @@ bytes_subscript(PyObject *op, PyObject* item)
         Py_ssize_t i = PyNumber_AsSsize_t(item, PyExc_IndexError);
         if (i == -1 && PyErr_Occurred())
             return NULL;
-        if (i < 0)
-            i += PyBytes_GET_SIZE(self);
-        if (i < 0 || i >= PyBytes_GET_SIZE(self)) {
-            PyErr_SetString(PyExc_IndexError,
-                            "index out of range");
-            return NULL;
-        }
-        return _PyLong_FromUnsignedChar((unsigned char)self->ob_sval[i]);
+        return _PyBytes_SubscriptIndex(op, i);
     }
     else if (PySlice_Check(item)) {
         Py_ssize_t start, stop, step, slicelength, i;

Python/specialize.c

@@ -2,7 +2,7 @@
 
 #include "opcode.h"
 
-#include "pycore_bytesobject.h"   // _PyBytes_Concat
+#include "pycore_bytesobject.h"   // _PyBytes_Concat(), _PyBytes_SubscriptIndex()
 #include "pycore_code.h"
 #include "pycore_critical_section.h"
 #include "pycore_descrobject.h"   // _PyMethodWrapper_Type
@@ -2193,6 +2193,26 @@ BITWISE_LONGS_ACTION(compactlongs_and, &)
 BITWISE_LONGS_ACTION(compactlongs_xor, ^)
 #undef BITWISE_LONGS_ACTION
 
+/* bytes subscripting */
+
+static inline int
+bytes_compactlong_index_in_bounds_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 offset;
+    return _PyBytes_OffsetFromIndex(lhs, index, &offset);
+}
+
+static PyObject *
+bytes_compactlong_subscr(PyObject *lhs, PyObject *rhs)
+{
+    Py_ssize_t index = _PyLong_CompactValue((PyLongObject *)rhs);
+    return _PyBytes_SubscriptIndex(lhs, index);
+}
+
 /* float-long */
 
 static inline int
@@ -2273,6 +2293,10 @@ 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_index_in_bounds_guard, bytes_compactlong_subscr, &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},