[mypyc] Add inline primitives for bytes.__getitem__

Author: JukkaL

mypyc/irbuild/specialize.py

@@ -94,7 +94,13 @@
     join_formatted_strings,
     tokenizer_format_call,
 )
-from mypyc.primitives.bytes_ops import isinstance_bytearray, isinstance_bytes
+from mypyc.primitives.bytes_ops import (
+    bytes_adjust_index_op,
+    bytes_get_item_unsafe_op,
+    bytes_range_check_op,
+    isinstance_bytearray,
+    isinstance_bytes,
+)
 from mypyc.primitives.dict_ops import (
     dict_items_op,
     dict_keys_op,
@@ -1299,3 +1305,54 @@ def translate_bytes_writer_set_item(
     )
 
     return builder.none()
+
+
+@specialize_dunder("__getitem__", bytes_rprimitive)
+def translate_bytes_get_item(
+    builder: IRBuilder, base_expr: Expression, args: list[Expression], ctx_expr: Expression
+) -> Value | None:
+    """Optimized bytes.__getitem__ implementation with bounds checking."""
+    # Check that we have exactly one argument
+    if len(args) != 1:
+        return None
+
+    # Get the bytes object
+    obj = builder.accept(base_expr)
+
+    # Get the index argument
+    index = builder.accept(args[0])
+
+    # Only use the optimized version for i64 index (requires experimental mode)
+    if not is_int64_rprimitive(index.type):
+        return None
+
+    # Adjust the index (handle negative indices)
+    adjusted_index = builder.primitive_op(
+        bytes_adjust_index_op, [obj, index], ctx_expr.line
+    )
+
+    # Check if the adjusted index is in valid range
+    range_check = builder.primitive_op(
+        bytes_range_check_op, [obj, adjusted_index], ctx_expr.line
+    )
+
+    # Create blocks for branching
+    valid_block = BasicBlock()
+    invalid_block = BasicBlock()
+
+    builder.add_bool_branch(range_check, valid_block, invalid_block)
+
+    # Handle invalid index - raise IndexError
+    builder.activate_block(invalid_block)
+    builder.add(
+        RaiseStandardError(RaiseStandardError.INDEX_ERROR, "index out of range", ctx_expr.line)
+    )
+    builder.add(Unreachable())
+
+    # Handle valid index - get the item
+    builder.activate_block(valid_block)
+    result = builder.primitive_op(
+        bytes_get_item_unsafe_op, [obj, adjusted_index], ctx_expr.line
+    )
+
+    return result

mypyc/lib-rt/bytes_extra_ops.h

@@ -2,9 +2,30 @@
 #define MYPYC_BYTES_EXTRA_OPS_H
 
 #include <Python.h>
+#include <stdint.h>
 #include "CPy.h"
 
 // Optimized bytes translate operation
 PyObject *CPyBytes_Translate(PyObject *bytes, PyObject *table);
 
+// Optimized bytes.__getitem__ operations
+
+// If index is negative, convert to non-negative index (no range checking)
+static inline int64_t CPyBytes_AdjustIndex(PyObject *obj, int64_t index) {
+    if (index < 0) {
+        return index + Py_SIZE(obj);
+    }
+    return index;
+}
+
+// Check if index is in valid range [0, len)
+static inline bool CPyBytes_RangeCheck(PyObject *obj, int64_t index) {
+    return index >= 0 && index < Py_SIZE(obj);
+}
+
+// Get byte at index (no bounds checking) - returns as CPyTagged
+static inline CPyTagged CPyBytes_GetItemUnsafe(PyObject *obj, int64_t index) {
+    return ((CPyTagged)(uint8_t)(PyBytes_AS_STRING(obj))[index]) << 1;
+}
+
 #endif

mypyc/primitives/bytes_ops.py

@@ -12,15 +12,18 @@
     c_int_rprimitive,
     c_pyssize_t_rprimitive,
     dict_rprimitive,
+    int64_rprimitive,
     int_rprimitive,
     list_rprimitive,
     object_rprimitive,
     str_rprimitive,
+    uint8_rprimitive,
 )
 from mypyc.primitives.registry import (
     ERR_NEG_INT,
     binary_op,
     custom_op,
+    custom_primitive_op,
     function_op,
     load_address_op,
     method_op,
@@ -166,3 +169,38 @@
     c_function_name="CPyBytes_Ord",
     error_kind=ERR_MAGIC,
 )
+
+# Optimized bytes.__getitem__ operations
+
+# bytes index adjustment - convert negative index to positive
+bytes_adjust_index_op = custom_primitive_op(
+    name="bytes_adjust_index",
+    arg_types=[bytes_rprimitive, int64_rprimitive],
+    return_type=int64_rprimitive,
+    c_function_name="CPyBytes_AdjustIndex",
+    error_kind=ERR_NEVER,
+    experimental=True,
+    dependencies=[BYTES_EXTRA_OPS],
+)
+
+# bytes range check - check if index is in valid range
+bytes_range_check_op = custom_primitive_op(
+    name="bytes_range_check",
+    arg_types=[bytes_rprimitive, int64_rprimitive],
+    return_type=bool_rprimitive,
+    c_function_name="CPyBytes_RangeCheck",
+    error_kind=ERR_NEVER,
+    experimental=True,
+    dependencies=[BYTES_EXTRA_OPS],
+)
+
+# bytes.__getitem__() - get byte at index (no bounds checking)
+bytes_get_item_unsafe_op = custom_primitive_op(
+    name="bytes_get_item_unsafe",
+    arg_types=[bytes_rprimitive, int64_rprimitive],
+    return_type=int_rprimitive,
+    c_function_name="CPyBytes_GetItemUnsafe",
+    error_kind=ERR_NEVER,
+    experimental=True,
+    dependencies=[BYTES_EXTRA_OPS],
+)

mypyc/test-data/irbuild-bytes.test

@@ -113,6 +113,28 @@ L0:
     r0 = CPyBytes_GetItem(a, i)
     return r0
 
+[case testBytesIndex_experimental_64bit]
+from mypy_extensions import i64
+
+def f(a: bytes, i: i64) -> int:
+    return a[i]
+[out]
+def f(a, i):
+    a :: bytes
+    i, r0 :: i64
+    r1, r2 :: bool
+    r3 :: int
+L0:
+    r0 = CPyBytes_AdjustIndex(a, i)
+    r1 = CPyBytes_RangeCheck(a, r0)
+    if r1 goto L2 else goto L1 :: bool
+L1:
+    r2 = raise IndexError('index out of range')
+    unreachable
+L2:
+    r3 = CPyBytes_GetItemUnsafe(a, r0)
+    return r3
+
 [case testBytesConcat]
 def f(a: bytes, b: bytes) -> bytes:
     return a + b