[mypyc] specialise bytearray(bytes()[i:j]) (#21901)

Fixes mypyc/mypyc#1216

Local compiled benchmarks showed approximately 1.8 to 2.1x speedup for 1
KiB and 1 MiB slices.

Authored-by: Codex

Author: hauntsaninja

mypyc/irbuild/specialize.py

@@ -30,6 +30,7 @@
     MemberExpr,
     NameExpr,
     RefExpr,
+    SliceExpr,
     StrExpr,
     SuperExpr,
     TupleExpr,
@@ -41,6 +42,7 @@
     Call,
     Extend,
     Integer,
+    LoadErrorValue,
     PrimitiveDescription,
     RaiseStandardError,
     Register,
@@ -64,7 +66,9 @@
     int32_rprimitive,
     int64_rprimitive,
     int_rprimitive,
+    is_any_int,
     is_bool_rprimitive,
+    is_bytes_rprimitive,
     is_dict_rprimitive,
     is_fixed_width_rtype,
     is_float_rprimitive,
@@ -108,7 +112,7 @@
     vec_to_list,
     vec_to_tuple,
 )
-from mypyc.primitives.bytearray_ops import isinstance_bytearray
+from mypyc.primitives.bytearray_ops import bytearray_from_bytes_slice_op, isinstance_bytearray
 from mypyc.primitives.bytes_ops import (
     bytes_adjust_index_op,
     bytes_get_item_unsafe_op,
@@ -348,6 +352,38 @@ def translate_vec_to_list(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -
     return None
 
 
+@specialize_function("builtins.bytearray")
+def translate_bytearray_from_bytes_slice(
+    builder: IRBuilder, expr: CallExpr, callee: RefExpr
+) -> Value | None:
+    """Construct a bytearray from a bytes slice without an intermediate copy."""
+    if len(expr.args) != 1 or expr.arg_kinds != [ARG_POS]:
+        return None
+    arg = expr.args[0]
+    if not isinstance(arg, IndexExpr) or not is_bytes_rprimitive(builder.node_type(arg.base)):
+        return None
+    index = arg.index
+    if (
+        not isinstance(index, SliceExpr)
+        or index.stride is not None
+        or (index.begin_index is not None and not is_any_int(builder.node_type(index.begin_index)))
+        or (index.end_index is not None and not is_any_int(builder.node_type(index.end_index)))
+    ):
+        return None
+
+    obj = builder.accept(arg.base)
+    # Use the default-argument sentinel so subclass slicing still receives None.
+    if index.begin_index is None:
+        start = builder.add(LoadErrorValue(int_rprimitive, is_borrowed=True))
+    else:
+        start = builder.accept(index.begin_index)
+    if index.end_index is None:
+        end = builder.add(LoadErrorValue(int_rprimitive, is_borrowed=True))
+    else:
+        end = builder.accept(index.end_index)
+    return builder.primitive_op(bytearray_from_bytes_slice_op, [obj, start, end], expr.line)
+
+
 @specialize_function("builtins.list")
 def dict_methods_fast_path(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
     """Specialize a common case when list() is called on a dictionary

mypyc/lib-rt/CPy.h

@@ -658,6 +658,7 @@ PyObject *CPyObject_GetAttr3(PyObject *v, PyObject *name, PyObject *defl);
 PyObject *CPyIter_Next(PyObject *iter);
 PyObject *CPyNumber_Power(PyObject *base, PyObject *index);
 PyObject *CPyNumber_InPlacePower(PyObject *base, PyObject *index);
+// An omitted slice bound is represented by CPY_INT_TAG.
 PyObject *CPyObject_GetSlice(PyObject *obj, CPyTagged start, CPyTagged end);
 
 

mypyc/lib-rt/bytearray_extra_ops.c

@@ -3,3 +3,31 @@
 PyObject *CPyByteArray_New(void) {
     return PyByteArray_FromStringAndSize(NULL, 0);
 }
+
+PyObject *CPyByteArray_FromBytesSlice(PyObject *obj, CPyTagged start, CPyTagged end) {
+    if (PyBytes_CheckExact(obj)
+            && (start == CPY_INT_TAG || CPyTagged_CheckShort(start))
+            && (end == CPY_INT_TAG || CPyTagged_CheckShort(end))) {
+        Py_ssize_t size = PyBytes_GET_SIZE(obj);
+        Py_ssize_t startn = start == CPY_INT_TAG ? 0 : CPyTagged_ShortAsSsize_t(start);
+        Py_ssize_t endn = end == CPY_INT_TAG ? size : CPyTagged_ShortAsSsize_t(end);
+        if (startn < 0) {
+            startn += size;
+        }
+        if (endn < 0) {
+            endn += size;
+        }
+        if (0 <= startn && startn <= endn && endn <= size) {
+            return PyByteArray_FromStringAndSize(PyBytes_AS_STRING(obj) + startn, endn - startn);
+        }
+    }
+
+    // Preserve general slice semantics, including bytes subclass overrides.
+    PyObject *slice = CPyObject_GetSlice(obj, start, end);
+    if (slice == NULL) {
+        return NULL;
+    }
+    PyObject *result = PyByteArray_FromObject(slice);
+    Py_DECREF(slice);
+    return result;
+}

mypyc/lib-rt/bytearray_extra_ops.h

@@ -7,4 +7,8 @@
 // Construct empty bytearray
 PyObject *CPyByteArray_New(void);
 
+// Construct a bytearray from a bytes slice, avoiding an intermediate bytes object.
+// An omitted bound is represented by CPY_INT_TAG.
+PyObject *CPyByteArray_FromBytesSlice(PyObject *obj, CPyTagged start, CPyTagged end);
+
 #endif

mypyc/lib-rt/generic_ops.c

@@ -47,9 +47,11 @@ PyObject *CPyNumber_InPlacePower(PyObject *base, PyObject *index)
 }
 
 PyObject *CPyObject_GetSlice(PyObject *obj, CPyTagged start, CPyTagged end) {
-    PyObject *start_obj = CPyTagged_AsObject(start);
-    PyObject *end_obj = CPyTagged_AsObject(end);
+    PyObject *start_obj = start == CPY_INT_TAG ? Py_NewRef(Py_None) : CPyTagged_AsObject(start);
+    PyObject *end_obj = end == CPY_INT_TAG ? Py_NewRef(Py_None) : CPyTagged_AsObject(end);
     if (unlikely(start_obj == NULL || end_obj == NULL)) {
+        Py_XDECREF(start_obj);
+        Py_XDECREF(end_obj);
         return NULL;
     }
     PyObject *slice = PySlice_New(start_obj, end_obj, NULL);

mypyc/primitives/bytearray_ops.py

@@ -9,7 +9,13 @@
 
 from mypyc.ir.deps import BYTEARRAY_EXTRA_OPS
 from mypyc.ir.ops import ERR_MAGIC, ERR_NEVER
-from mypyc.ir.rtypes import bit_rprimitive, bytearray_rprimitive, object_rprimitive
+from mypyc.ir.rtypes import (
+    bit_rprimitive,
+    bytearray_rprimitive,
+    bytes_rprimitive,
+    int_rprimitive,
+    object_rprimitive,
+)
 from mypyc.primitives.registry import custom_primitive_op, function_op, load_address_op
 
 # Get the 'bytearray' type object.
@@ -24,6 +30,17 @@
     error_kind=ERR_MAGIC,
 )
 
+# bytearray(bytes[start:end])
+# Omitted bounds use the tagged integer error value.
+bytearray_from_bytes_slice_op = custom_primitive_op(
+    name="bytearray_from_bytes_slice",
+    arg_types=[bytes_rprimitive, int_rprimitive, int_rprimitive],
+    return_type=bytearray_rprimitive,
+    c_function_name="CPyByteArray_FromBytesSlice",
+    error_kind=ERR_MAGIC,
+    dependencies=[BYTEARRAY_EXTRA_OPS],
+)
+
 # bytearray() -- construct empty bytearray
 function_op(
     name="builtins.bytearray",

mypyc/test-data/irbuild-bytes.test

@@ -99,6 +99,106 @@ L0:
     r0 = CPyBytes_GetSlice(a, start, end)
     return r0
 
+[case testBytearrayFromBytesSlice]
+def f(a: bytes, start: int, end: int) -> bytearray:
+    return bytearray(a[start:end])
+
+def from_start(a: bytes, start: int) -> bytearray:
+    return bytearray(a[start:])
+
+def to_end(a: bytes, end: int) -> bytearray:
+    return bytearray(a[:end])
+
+def full(a: bytes) -> bytearray:
+    return bytearray(a[:])
+[out]
+def f(a, start, end):
+    a :: bytes
+    start, end :: int
+    r0 :: bytearray
+L0:
+    r0 = CPyByteArray_FromBytesSlice(a, start, end)
+    return r0
+def from_start(a, start):
+    a :: bytes
+    start, r0 :: int
+    r1 :: bytearray
+L0:
+    r0 = <error> :: int
+    r1 = CPyByteArray_FromBytesSlice(a, start, r0)
+    return r1
+def to_end(a, end):
+    a :: bytes
+    end, r0 :: int
+    r1 :: bytearray
+L0:
+    r0 = <error> :: int
+    r1 = CPyByteArray_FromBytesSlice(a, r0, end)
+    return r1
+def full(a):
+    a :: bytes
+    r0, r1 :: int
+    r2 :: bytearray
+L0:
+    r0 = <error> :: int
+    r1 = <error> :: int
+    r2 = CPyByteArray_FromBytesSlice(a, r0, r1)
+    return r2
+
+[case testBytearrayFromBytesSliceFixedWidth_64bit]
+from mypy_extensions import i32, i64
+
+def f(a: bytes, start: i32, end: i64) -> bytearray:
+    return bytearray(a[start:end])
+[out]
+def f(a, start, end):
+    a :: bytes
+    start :: i32
+    end :: i64
+    r0 :: native_int
+    r1 :: int
+    r2, r3 :: bit
+    r4, r5, r6 :: int
+    r7 :: bytearray
+L0:
+    r0 = extend signed start: i32 to native_int
+    r1 = r0 << 1
+    r2 = end <= 4611686018427387903 :: signed
+    if r2 goto L1 else goto L2 :: bool
+L1:
+    r3 = end >= -4611686018427387904 :: signed
+    if r3 goto L3 else goto L2 :: bool
+L2:
+    r4 = CPyTagged_FromInt64(end)
+    r5 = r4
+    goto L4
+L3:
+    r6 = end << 1
+    r5 = r6
+L4:
+    r7 = CPyByteArray_FromBytesSlice(a, r1, r5)
+    return r7
+
+[case testBytearrayFromBytesSliceWithStep]
+def f(a: bytes, start: int, end: int, step: int) -> bytearray:
+    return bytearray(a[start:end:step])
+[out]
+def f(a, start, end, step):
+    a :: bytes
+    start, end, step :: int
+    r0, r1, r2, r3, r4 :: object
+    r5 :: bytes
+    r6 :: bytearray
+L0:
+    r0 = box(int, start)
+    r1 = box(int, end)
+    r2 = box(int, step)
+    r3 = PySlice_New(r0, r1, r2)
+    r4 = PyObject_GetItem(a, r3)
+    r5 = cast(bytes, r4)
+    r6 = PyByteArray_FromObject(r5)
+    return r6
+
 [case testBytesIndex]
 from mypy_extensions import i64