Merge branch 'master' into to-bytes

mypyc/irbuild/specialize.py

@@ -72,7 +72,9 @@
     is_int64_rprimitive,
     is_int_rprimitive,
     is_list_rprimitive,
+    is_sequence_rprimitive,
     is_str_rprimitive,
+    is_tagged,
     is_uint8_rprimitive,
     list_rprimitive,
     object_rprimitive,
@@ -117,6 +119,14 @@
     int_to_little_endian_op,
     isinstance_int,
 )
+from mypyc.primitives.generic_ops import generic_setattr, setup_object
+from mypyc.primitives.int_ops import isinstance_int
+from mypyc.primitives.librt_strings_ops import (
+    bytes_writer_adjust_index_op,
+    bytes_writer_get_item_unsafe_op,
+    bytes_writer_range_check_op,
+    bytes_writer_set_item_unsafe_op,
+)
 from mypyc.primitives.list_ops import isinstance_list, new_list_set_item_op
 from mypyc.primitives.misc_ops import isinstance_bool
 from mypyc.primitives.set_ops import isinstance_frozenset, isinstance_set
@@ -1273,4 +1283,148 @@ def specialize_int_to_bytes(builder: IRBuilder, expr: CallExpr, callee: RefExpr)
     byteorder_arg = builder.accept(byteorder_expr)
     return builder.call_c(
         int_to_bytes_op, [self_arg, length_arg, byteorder_arg, signed_arg], expr.line
+
+
+def translate_getitem_with_bounds_check(
+    builder: IRBuilder,
+    base_expr: Expression,
+    args: list[Expression],
+    ctx_expr: Expression,
+    adjust_index_op: PrimitiveDescription,
+    range_check_op: PrimitiveDescription,
+    get_item_unsafe_op: PrimitiveDescription,
+) -> Value | None:
+    """Shared helper for optimized __getitem__ with bounds checking.
+
+    This implements the common pattern of:
+    1. Adjusting negative indices
+    2. Checking if index is in valid range
+    3. Raising IndexError if out of range
+    4. Getting the item if in range
+
+    Args:
+        builder: The IR builder
+        base_expr: The base object expression
+        args: The arguments to __getitem__ (should be length 1)
+        ctx_expr: The context expression for line numbers
+        adjust_index_op: Primitive op to adjust negative indices
+        range_check_op: Primitive op to check if index is in valid range
+        get_item_unsafe_op: Primitive op to get item (no bounds checking)
+
+    Returns:
+        The result value, or None if optimization doesn't apply
+    """
+    # Check that we have exactly one argument
+    if len(args) != 1:
+        return None
+
+    # Get the object
+    obj = builder.accept(base_expr)
+
+    # Get the index argument
+    index = builder.accept(args[0])
+
+    # Adjust the index (handle negative indices)
+    adjusted_index = builder.primitive_op(adjust_index_op, [obj, index], ctx_expr.line)
+
+    # Check if the adjusted index is in valid range
+    range_check = builder.primitive_op(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(get_item_unsafe_op, [obj, adjusted_index], ctx_expr.line)
+
+    return result
+
+
+@specialize_dunder("__getitem__", bytes_writer_rprimitive)
+def translate_bytes_writer_get_item(
+    builder: IRBuilder, base_expr: Expression, args: list[Expression], ctx_expr: Expression
+) -> Value | None:
+    """Optimized BytesWriter.__getitem__ implementation with bounds checking."""
+    return translate_getitem_with_bounds_check(
+        builder,
+        base_expr,
+        args,
+        ctx_expr,
+        bytes_writer_adjust_index_op,
+        bytes_writer_range_check_op,
+        bytes_writer_get_item_unsafe_op,
+    )
+
+
+@specialize_dunder("__setitem__", bytes_writer_rprimitive)
+def translate_bytes_writer_set_item(
+    builder: IRBuilder, base_expr: Expression, args: list[Expression], ctx_expr: Expression
+) -> Value | None:
+    """Optimized BytesWriter.__setitem__ implementation with bounds checking."""
+    # Check that we have exactly two arguments (index and value)
+    if len(args) != 2:
+        return None
+
+    # Get the BytesWriter object
+    obj = builder.accept(base_expr)
+
+    # Get the index and value arguments
+    index = builder.accept(args[0])
+    value = builder.accept(args[1])
+
+    # Adjust the index (handle negative indices)
+    adjusted_index = builder.primitive_op(
+        bytes_writer_adjust_index_op, [obj, index], ctx_expr.line
+    )
+
+    # Check if the adjusted index is in valid range
+    range_check = builder.primitive_op(
+        bytes_writer_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 - set the item
+    builder.activate_block(valid_block)
+    builder.primitive_op(
+        bytes_writer_set_item_unsafe_op, [obj, adjusted_index, value], ctx_expr.line
+    )
+
+    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."""
+    return translate_getitem_with_bounds_check(
+        builder,
+        base_expr,
+        args,
+        ctx_expr,
+        bytes_adjust_index_op,
+        bytes_range_check_op,
+        bytes_get_item_unsafe_op,
     )