[mypyc] Add primitives for bytearray len, indexing, item assignment and slicing

mypyc/doc/bytearray_operations.rst

@@ -11,3 +11,21 @@ Construction
 
 * ``bytearray()``
 * ``bytearray(x)``
+
+Operators
+---------
+
+* Indexing (``b[n]``)
+* Slicing (``b[n:m]``, ``b[n:]``, ``b[:m]``)
+
+Statements
+----------
+
+Set item by integer index:
+
+* ``b[n] = x``
+
+Functions
+---------
+
+* ``len(b: bytearray)``

mypyc/irbuild/expression.py

@@ -146,6 +146,7 @@
     vec_item_type,
     vec_slice,
 )
+from mypyc.primitives.bytearray_ops import bytearray_slice_op
 from mypyc.primitives.bytes_ops import bytes_slice_op
 from mypyc.primitives.dict_ops import dict_get_item_op, dict_new_op, exact_dict_set_item_op
 from mypyc.primitives.float_ops import (
@@ -1084,6 +1085,8 @@ def try_gen_slice_op(builder: IRBuilder, base: Value, index: SliceExpr) -> Value
             end = builder.load_int(MAX_SHORT_INT)
         if isinstance(base.type, RVec):
             return vec_slice(builder.builder, base, begin, end, index.line)
+        if is_bytearray_rprimitive(base.type):
+            return builder.primitive_op(bytearray_slice_op, [base, begin, end], index.line)
         candidates = [list_slice_op, tuple_slice_op, str_slice_op, bytes_slice_op]
         return builder.builder.matching_call_c(candidates, [base, begin, end], index.line)
 

mypyc/irbuild/ll_builder.py

@@ -102,6 +102,7 @@
     int64_rprimitive,
     int_rprimitive,
     is_bool_or_bit_rprimitive,
+    is_bytearray_rprimitive,
     is_bytes_rprimitive,
     is_c_py_ssize_t_rprimitive,
     is_dict_rprimitive,
@@ -2893,7 +2894,12 @@ def builtin_len(self, val: Value, line: int, use_pyssize_t: bool = False) -> Val
         """
         typ = val.type
         size_value = None
-        if is_list_rprimitive(typ) or is_tuple_rprimitive(typ) or is_bytes_rprimitive(typ):
+        if (
+            is_list_rprimitive(typ)
+            or is_tuple_rprimitive(typ)
+            or is_bytes_rprimitive(typ)
+            or is_bytearray_rprimitive(typ)
+        ):
             size_value = self.primitive_op(var_object_size, [val], line)
         elif is_set_rprimitive(typ) or is_frozenset_rprimitive(typ):
             elem_address = self.add(GetElementPtr(val, PySetObject, "used"))

mypyc/lib-rt/bytearray_extra_ops.c

@@ -31,3 +31,131 @@ PyObject *CPyByteArray_FromBytesSlice(PyObject *obj, CPyTagged start, CPyTagged
     Py_DECREF(slice);
     return result;
 }
+
+// On free-threaded builds, the item array of a bytearray may be resized or freed
+// concurrently, so it must only be accessed while holding the object's lock.
+#ifdef Py_GIL_DISABLED
+#define CPY_BYTEARRAY_LOCK(obj) Py_BEGIN_CRITICAL_SECTION(obj)
+#define CPY_BYTEARRAY_UNLOCK() Py_END_CRITICAL_SECTION()
+#else
+#define CPY_BYTEARRAY_LOCK(obj) {
+#define CPY_BYTEARRAY_UNLOCK() }
+#endif
+
+CPyTagged CPyByteArray_GetItem(PyObject *obj, CPyTagged index) {
+    if (likely(PyByteArray_CheckExact(obj) && CPyTagged_CheckShort(index))) {
+        Py_ssize_t n = CPyTagged_ShortAsSsize_t(index);
+        CPyTagged result = CPY_INT_TAG;
+        CPY_BYTEARRAY_LOCK(obj);
+        Py_ssize_t size = PyByteArray_GET_SIZE(obj);
+        if (n < 0) {
+            n += size;
+        }
+        if (n < 0 || n >= size) {
+            PyErr_SetString(PyExc_IndexError, "bytearray index out of range");
+        } else {
+            result = (CPyTagged)(unsigned char)PyByteArray_AS_STRING(obj)[n] << 1;
+        }
+        CPY_BYTEARRAY_UNLOCK();
+        return result;
+    }
+
+    // Generic path for subclasses and large indexes
+    PyObject *index_obj = CPyTagged_AsObject(index);
+    if (index_obj == NULL) {
+        return CPY_INT_TAG;
+    }
+    PyObject *item = PyObject_GetItem(obj, index_obj);
+    Py_DECREF(index_obj);
+    if (item == NULL) {
+        return CPY_INT_TAG;
+    }
+    if (unlikely(!PyLong_Check(item))) {
+        CPy_TypeError("int", item);
+        Py_DECREF(item);
+        return CPY_INT_TAG;
+    }
+    return CPyTagged_StealFromObject(item);
+}
+
+bool CPyByteArray_SetItem(PyObject *obj, CPyTagged index, CPyTagged value) {
+    if (likely(PyByteArray_CheckExact(obj) && CPyTagged_CheckShort(index))) {
+        Py_ssize_t n = CPyTagged_ShortAsSsize_t(index);
+        // Like CPython, check the value before the index
+        Py_ssize_t v = CPyTagged_CheckShort(value) ? CPyTagged_ShortAsSsize_t(value) : -1;
+        if (v < 0 || v > 255) {
+            PyErr_SetString(PyExc_ValueError, "byte must be in range(0, 256)");
+            return false;
+        }
+        bool ok = false;
+        CPY_BYTEARRAY_LOCK(obj);
+        Py_ssize_t size = PyByteArray_GET_SIZE(obj);
+        if (n < 0) {
+            n += size;
+        }
+        if (n < 0 || n >= size) {
+            PyErr_SetString(PyExc_IndexError, "bytearray index out of range");
+        } else {
+            PyByteArray_AS_STRING(obj)[n] = (char)v;
+            ok = true;
+        }
+        CPY_BYTEARRAY_UNLOCK();
+        return ok;
+    }
+
+    // Generic path for subclasses and large indexes
+    PyObject *index_obj = CPyTagged_AsObject(index);
+    if (index_obj == NULL) {
+        return false;
+    }
+    PyObject *value_obj = CPyTagged_AsObject(value);
+    if (value_obj == NULL) {
+        Py_DECREF(index_obj);
+        return false;
+    }
+    int res = PyObject_SetItem(obj, index_obj, value_obj);
+    Py_DECREF(index_obj);
+    Py_DECREF(value_obj);
+    return res >= 0;
+}
+
+PyObject *CPyByteArray_GetSlice(PyObject *obj, CPyTagged start, CPyTagged end) {
+    if (likely(PyByteArray_CheckExact(obj)
+               && CPyTagged_CheckShort(start) && CPyTagged_CheckShort(end))) {
+        Py_ssize_t startn = CPyTagged_ShortAsSsize_t(start);
+        Py_ssize_t endn = CPyTagged_ShortAsSsize_t(end);
+        PyObject *result;
+        CPY_BYTEARRAY_LOCK(obj);
+        Py_ssize_t size = PyByteArray_GET_SIZE(obj);
+        if (startn < 0) {
+            startn += size;
+            if (startn < 0) {
+                startn = 0;
+            }
+        } else if (startn > size) {
+            startn = size;
+        }
+        if (endn < 0) {
+            endn += size;
+            if (endn < 0) {
+                endn = 0;
+            }
+        } else if (endn > size) {
+            endn = size;
+        }
+        Py_ssize_t len = endn > startn ? endn - startn : 0;
+        result = PyByteArray_FromStringAndSize(PyByteArray_AS_STRING(obj) + startn, len);
+        CPY_BYTEARRAY_UNLOCK();
+        return result;
+    }
+
+    // Generic path for subclasses and large indexes
+    PyObject *result = CPyObject_GetSlice(obj, start, end);
+    if (result != NULL && unlikely(!PyByteArray_Check(result))) {
+        // A subclass may override __getitem__ to return anything
+        CPy_TypeError("bytearray", result);
+        Py_DECREF(result);
+        return NULL;
+    }
+    return result;
+}

mypyc/lib-rt/bytearray_extra_ops.h

@@ -11,4 +11,13 @@ PyObject *CPyByteArray_New(void);
 // An omitted bound is represented by CPY_INT_TAG.
 PyObject *CPyByteArray_FromBytesSlice(PyObject *obj, CPyTagged start, CPyTagged end);
 
+// bytearray[index]
+CPyTagged CPyByteArray_GetItem(PyObject *obj, CPyTagged index);
+
+// bytearray[index] = value
+bool CPyByteArray_SetItem(PyObject *obj, CPyTagged index, CPyTagged value);
+
+// bytearray[start:end]
+PyObject *CPyByteArray_GetSlice(PyObject *obj, CPyTagged start, CPyTagged end);
+
 #endif

mypyc/primitives/bytearray_ops.py

@@ -8,15 +8,15 @@
 from __future__ import annotations
 
 from mypyc.ir.deps import BYTEARRAY_EXTRA_OPS
-from mypyc.ir.ops import ERR_MAGIC, ERR_NEVER
+from mypyc.ir.ops import ERR_FALSE, ERR_MAGIC, ERR_NEVER
 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
+from mypyc.primitives.registry import custom_primitive_op, function_op, load_address_op, method_op
 
 # Get the 'bytearray' type object.
 load_address_op(name="builtins.bytearray", type=object_rprimitive, src="PyByteArray_Type")
@@ -59,3 +59,33 @@
     c_function_name="PyByteArray_Check",
     error_kind=ERR_NEVER,
 )
+
+# bytearray[index]
+method_op(
+    name="__getitem__",
+    arg_types=[bytearray_rprimitive, int_rprimitive],
+    return_type=int_rprimitive,
+    c_function_name="CPyByteArray_GetItem",
+    error_kind=ERR_MAGIC,
+    dependencies=[BYTEARRAY_EXTRA_OPS],
+)
+
+# bytearray[index] = value
+method_op(
+    name="__setitem__",
+    arg_types=[bytearray_rprimitive, int_rprimitive, int_rprimitive],
+    return_type=bit_rprimitive,
+    c_function_name="CPyByteArray_SetItem",
+    error_kind=ERR_FALSE,
+    dependencies=[BYTEARRAY_EXTRA_OPS],
+)
+
+# bytearray[begin:end]
+bytearray_slice_op = custom_primitive_op(
+    name="bytearray_slice",
+    arg_types=[bytearray_rprimitive, int_rprimitive, int_rprimitive],
+    return_type=bytearray_rprimitive,
+    c_function_name="CPyByteArray_GetSlice",
+    error_kind=ERR_MAGIC,
+    dependencies=[BYTEARRAY_EXTRA_OPS],
+)

mypyc/primitives/bytes_ops.py

@@ -94,7 +94,6 @@
 )
 
 # bytes[index]
-# bytearray[index]
 method_op(
     name="__getitem__",
     arg_types=[bytes_rprimitive, int_rprimitive],

mypyc/test-data/irbuild-bytes.test

@@ -257,6 +257,69 @@ L0:
     r1 = r0 << 1
     return r1
 
+[case testBytearrayLen]
+def f(b: bytearray) -> int:
+    return len(b)
+[out]
+def f(b):
+    b :: bytearray
+    r0 :: native_int
+    r1 :: short_int
+L0:
+    r0 = var_object_size b
+    r1 = r0 << 1
+    return r1
+
+[case testBytearrayGetItem]
+def f(b: bytearray, i: int) -> int:
+    return b[i]
+[out]
+def f(b, i):
+    b :: bytearray
+    i, r0 :: int
+L0:
+    r0 = CPyByteArray_GetItem(b, i)
+    return r0
+
+[case testBytearraySetItem]
+def f(b: bytearray, i: int, x: int) -> None:
+    b[i] = x
+[out]
+def f(b, i, x):
+    b :: bytearray
+    i, x :: int
+    r0 :: bit
+L0:
+    r0 = CPyByteArray_SetItem(b, i, x)
+    return 1
+
+[case testBytearraySlice]
+def f(b: bytearray, i: int, j: int) -> bytearray:
+    return b[i:j]
+def g(b: bytearray, i: int, j: int, k: int) -> bytearray:
+    return b[i:j:k]
+[out]
+def f(b, i, j):
+    b :: bytearray
+    i, j :: int
+    r0 :: bytearray
+L0:
+    r0 = CPyByteArray_GetSlice(b, i, j)
+    return r0
+def g(b, i, j, k):
+    b :: bytearray
+    i, j, k :: int
+    r0, r1, r2, r3, r4 :: object
+    r5 :: bytearray
+L0:
+    r0 = box(int, i)
+    r1 = box(int, j)
+    r2 = box(int, k)
+    r3 = PySlice_New(r0, r1, r2)
+    r4 = PyObject_GetItem(b, r3)
+    r5 = cast(bytearray, r4)
+    return r5
+
 [case testBytesFormatting]
 def f(var: bytes, num: int) -> None:
     b1 = b'aaaa%bbbbb%s' % (var, var)