[mypyc] Add get/set item, len and truncate to BytesWriter (#20298)

Add support for `b[i]`, `b[i] = x`, `len(b)` and `b.truncate(n)` to
`librt.strings.BytesWriter`.

The implementations are still unoptimized. Get item and set item don't
have primitives yet, but I'll add them in a follow-up PR.

Author: JukkaL

mypy/typeshed/stubs/librt/librt/strings.pyi

@@ -1,7 +1,13 @@
 from typing import final
 
+from mypy_extensions import i64, u8
+
 @final
 class BytesWriter:
     def append(self, /, x: int) -> None: ...
     def write(self, /, b: bytes) -> None: ...
     def getvalue(self) -> bytes: ...
+    def truncate(self, /, size: i64) -> None: ...
+    def __len__(self) -> i64: ...
+    def __getitem__(self, /, i: i64) -> u8: ...
+    def __setitem__(self, /, i: i64, x: u8) -> None: ...

mypyc/irbuild/ll_builder.py

@@ -179,6 +179,7 @@
     ERR_NEG_INT,
     CFunctionDescription,
     binary_ops,
+    function_ops,
     method_call_ops,
     unary_ops,
 )
@@ -2489,7 +2490,11 @@ def builtin_len(self, val: Value, line: int, use_pyssize_t: bool = False) -> Val
             self.activate_block(ok)
             return length
 
-        # generic case
+        op = self.matching_primitive_op(function_ops["builtins.len"], [val], line)
+        if op is not None:
+            return op
+
+        # Fallback generic case
         if use_pyssize_t:
             return self.call_c(generic_ssize_t_len_op, [val], line)
         else:

mypyc/lib-rt/librt_strings.c

@@ -41,7 +41,7 @@ _grow_buffer(BytesWriterObject *data, Py_ssize_t n) {
     do {
         size *= 2;
     } while (target >= size);
-    if (old_size == WRITER_EMBEDDED_BUF_LEN) {
+    if (data->buf == data->data) {
         // Move from embedded buffer to heap-allocated buffer
         data->buf = PyMem_Malloc(size);
         if (data->buf != NULL) {
@@ -155,8 +155,65 @@ BytesWriter_getvalue(BytesWriterObject *self, PyObject *Py_UNUSED(ignored))
     return PyBytes_FromStringAndSize(self->buf, self->len);
 }
 
+static Py_ssize_t
+BytesWriter_length(BytesWriterObject *self)
+{
+    return self->len;
+}
+
+static PyObject*
+BytesWriter_item(BytesWriterObject *self, Py_ssize_t index)
+{
+    Py_ssize_t length = self->len;
+
+    // Check bounds
+    if (index < 0 || index >= length) {
+        PyErr_SetString(PyExc_IndexError, "BytesWriter index out of range");
+        return NULL;
+    }
+
+    // Return the byte at the given index as a Python int
+    return PyLong_FromLong((unsigned char)self->buf[index]);
+}
+
+static int
+BytesWriter_ass_item(BytesWriterObject *self, Py_ssize_t index, PyObject *value)
+{
+    Py_ssize_t length = self->len;
+
+    // Check bounds
+    if (index < 0 || index >= length) {
+        PyErr_SetString(PyExc_IndexError, "BytesWriter index out of range");
+        return -1;
+    }
+
+    // Check that value is not NULL (deletion not supported)
+    if (value == NULL) {
+        PyErr_SetString(PyExc_TypeError, "BytesWriter does not support item deletion");
+        return -1;
+    }
+
+    // Convert value to uint8
+    uint8_t byte_value = CPyLong_AsUInt8(value);
+    if (unlikely(byte_value == CPY_LL_UINT_ERROR && PyErr_Occurred())) {
+        CPy_TypeError("u8", value);
+        return -1;
+    }
+
+    // Assign the byte
+    self->buf[index] = (char)byte_value;
+    return 0;
+}
+
+static PySequenceMethods BytesWriter_as_sequence = {
+    .sq_length = (lenfunc)BytesWriter_length,
+    .sq_item = (ssizeargfunc)BytesWriter_item,
+    .sq_ass_item = (ssizeobjargproc)BytesWriter_ass_item,
+};
+
 static PyObject* BytesWriter_append(PyObject *self, PyObject *const *args, size_t nargs, PyObject *kwnames);
 static PyObject* BytesWriter_write(PyObject *self, PyObject *const *args, size_t nargs, PyObject *kwnames);
+static PyObject* BytesWriter_truncate(PyObject *self, PyObject *const *args, size_t nargs);
 
 static PyMethodDef BytesWriter_methods[] = {
     {"append", (PyCFunction) BytesWriter_append, METH_FASTCALL | METH_KEYWORDS,
@@ -168,6 +225,9 @@ static PyMethodDef BytesWriter_methods[] = {
     {"getvalue", (PyCFunction) BytesWriter_getvalue, METH_NOARGS,
      "Return the buffer content as bytes object"
     },
+    {"truncate", (PyCFunction) BytesWriter_truncate, METH_FASTCALL,
+     PyDoc_STR("Truncate the buffer to the specified size")
+    },
     {NULL}  /* Sentinel */
 };
 
@@ -182,6 +242,7 @@ static PyTypeObject BytesWriterType = {
     .tp_init = (initproc) BytesWriter_init,
     .tp_dealloc = (destructor) BytesWriter_dealloc,
     .tp_methods = BytesWriter_methods,
+    .tp_as_sequence = &BytesWriter_as_sequence,
     .tp_repr = (reprfunc)BytesWriter_repr,
 };
 
@@ -268,11 +329,68 @@ BytesWriter_append(PyObject *self, PyObject *const *args, size_t nargs, PyObject
     return Py_None;
 }
 
+static char
+BytesWriter_truncate_internal(PyObject *self, int64_t size) {
+    BytesWriterObject *writer = (BytesWriterObject *)self;
+    Py_ssize_t current_size = writer->len;
+
+    // Validate size is non-negative
+    if (size < 0) {
+        PyErr_SetString(PyExc_ValueError, "size must be non-negative");
+        return CPY_NONE_ERROR;
+    }
+
+    // Validate size doesn't exceed current size
+    if (size > current_size) {
+        PyErr_SetString(PyExc_ValueError, "size cannot be larger than current buffer size");
+        return CPY_NONE_ERROR;
+    }
+
+    writer->len = size;
+    return CPY_NONE;
+}
+
+static PyObject*
+BytesWriter_truncate(PyObject *self, PyObject *const *args, size_t nargs) {
+    if (unlikely(nargs != 1)) {
+        PyErr_Format(PyExc_TypeError,
+                     "truncate() takes exactly 1 argument (%zu given)", nargs);
+        return NULL;
+    }
+    if (!check_bytes_writer(self)) {
+        return NULL;
+    }
+
+    PyObject *size_obj = args[0];
+    int overflow;
+    long long size = PyLong_AsLongLongAndOverflow(size_obj, &overflow);
+
+    if (size == -1 && PyErr_Occurred()) {
+        return NULL;
+    }
+    if (overflow != 0) {
+        PyErr_SetString(PyExc_ValueError, "integer out of range");
+        return NULL;
+    }
+
+    if (unlikely(BytesWriter_truncate_internal(self, size) == CPY_NONE_ERROR)) {
+        return NULL;
+    }
+    Py_INCREF(Py_None);
+    return Py_None;
+}
+
 static PyTypeObject *
 BytesWriter_type_internal(void) {
     return &BytesWriterType;  // Return borrowed reference
 };
 
+static CPyTagged
+BytesWriter_len_internal(PyObject *self) {
+    BytesWriterObject *writer = (BytesWriterObject *)self;
+    return writer->len << 1;
+}
+
 static PyMethodDef librt_strings_module_methods[] = {
     {NULL, NULL, 0, NULL}
 };
@@ -311,6 +429,8 @@ librt_strings_module_exec(PyObject *m)
         (void *)BytesWriter_append_internal,
         (void *)BytesWriter_write_internal,
         (void *)BytesWriter_type_internal,
+        (void *)BytesWriter_len_internal,
+        (void *)BytesWriter_truncate_internal,
     };
     PyObject *c_api_object = PyCapsule_New((void *)librt_strings_api, "librt.strings._C_API", NULL);
     if (PyModule_Add(m, "_C_API", c_api_object) < 0) {

mypyc/lib-rt/librt_strings.h

@@ -20,11 +20,11 @@ import_librt_strings(void)
 // API version -- more recent versions must maintain backward compatibility, i.e.
 // we can add new features but not remove or change existing features (unless
 // ABI version is changed, but see the comment above).
- #define LIBRT_STRINGS_API_VERSION 0
+ #define LIBRT_STRINGS_API_VERSION 1
 
 // Number of functions in the capsule API. If you add a new function, also increase
 // LIBRT_STRINGS_API_VERSION.
-#define LIBRT_STRINGS_API_LEN 7
+#define LIBRT_STRINGS_API_LEN 9
 
 static void *LibRTStrings_API[LIBRT_STRINGS_API_LEN];
 
@@ -35,6 +35,8 @@ static void *LibRTStrings_API[LIBRT_STRINGS_API_LEN];
 #define LibRTStrings_BytesWriter_append_internal (*(char (*)(PyObject *source, uint8_t value)) LibRTStrings_API[4])
 #define LibRTStrings_BytesWriter_write_internal (*(char (*)(PyObject *source, PyObject *value)) LibRTStrings_API[5])
 #define LibRTStrings_BytesWriter_type_internal (*(PyTypeObject* (*)(void)) LibRTStrings_API[6])
+#define LibRTStrings_BytesWriter_len_internal (*(CPyTagged (*)(PyObject *self)) LibRTStrings_API[7])
+#define LibRTStrings_BytesWriter_truncate_internal (*(char (*)(PyObject *self, int64_t size)) LibRTStrings_API[8])
 
 static int
 import_librt_strings(void)

mypyc/primitives/librt_strings_ops.py

@@ -1,7 +1,14 @@
 from typing import Final
 
-from mypyc.ir.ops import ERR_MAGIC
-from mypyc.ir.rtypes import KNOWN_NATIVE_TYPES, bytes_rprimitive, none_rprimitive, uint8_rprimitive
+from mypyc.ir.ops import ERR_MAGIC, ERR_NEVER
+from mypyc.ir.rtypes import (
+    KNOWN_NATIVE_TYPES,
+    bytes_rprimitive,
+    int64_rprimitive,
+    none_rprimitive,
+    short_int_rprimitive,
+    uint8_rprimitive,
+)
 from mypyc.primitives.registry import function_op, method_op
 
 bytes_writer_rprimitive: Final = KNOWN_NATIVE_TYPES["librt.strings.BytesWriter"]
@@ -45,3 +52,21 @@
     experimental=True,
     capsule="librt.strings",
 )
+
+method_op(
+    name="truncate",
+    arg_types=[bytes_writer_rprimitive, int64_rprimitive],
+    return_type=none_rprimitive,
+    c_function_name="LibRTStrings_BytesWriter_truncate_internal",
+    error_kind=ERR_MAGIC,
+)
+
+function_op(
+    name="builtins.len",
+    arg_types=[bytes_writer_rprimitive],
+    return_type=short_int_rprimitive,
+    c_function_name="LibRTStrings_BytesWriter_len_internal",
+    error_kind=ERR_NEVER,
+    experimental=True,
+    capsule="librt.strings",
+)

mypyc/primitives/registry.py

@@ -7,20 +7,20 @@
 
 Example op definition:
 
-list_len_op = func_op(name='builtins.len',
-                      arg_types=[list_rprimitive],
-                      result_type=short_int_rprimitive,
-                      error_kind=ERR_NEVER,
-                      emit=emit_len)
+list_len_op = function_op(name='builtins.len',
+                          arg_types=[list_rprimitive],
+                          result_type=short_int_rprimitive,
+                          error_kind=ERR_NEVER,
+                          c_function_name="...")
 
 This op is automatically generated for calls to len() with a single
 list argument. The result type is short_int_rprimitive, and this
-never raises an exception (ERR_NEVER). The function emit_len is used
-to generate C for this op.  The op can also be manually generated using
-"list_len_op". Ops that are only generated automatically don't need to
+never raises an exception (ERR_NEVER). The function c_function_name is
+called when generating C for this op.  The op can also be manually generated
+using "list_len_op". Ops that are only generated automatically don't need to
 be assigned to a module attribute.
 
-Ops defined with custom_op are only explicitly generated in
+Ops defined with custom[_primitive]_op are only explicitly generated in
 mypyc.irbuild and won't be generated automatically. They are always
 assigned to a module attribute, as otherwise they won't be accessible.
 
@@ -306,6 +306,7 @@ def custom_primitive_op(
     steals: StealsDescription = False,
     is_borrowed: bool = False,
     is_pure: bool = False,
+    experimental: bool = False,
     capsule: str | None = None,
 ) -> PrimitiveDescription:
     """Define a primitive op that can't be automatically generated based on the AST.
@@ -328,7 +329,7 @@ def custom_primitive_op(
         extra_int_constants=extra_int_constants,
         priority=0,
         is_pure=is_pure,
-        experimental=False,
+        experimental=experimental,
         capsule=capsule,
     )
 

mypyc/test-data/irbuild-librt-strings.test

@@ -1,27 +1,43 @@
 [case testLibrtStrings_experimental]
 from librt.strings import BytesWriter
-from mypy_extensions import u8
+from mypy_extensions import u8, i64
 
-def bytes_builder_basics() -> bytes:
+def bytes_writer_basics() -> bytes:
     b = BytesWriter()
     x: u8 = 1
     b.append(x)
     b.write(b'foo')
+    n: i64 = 2
+    b.truncate(n)
     return b.getvalue()
+def bytes_writer_len(b: BytesWriter) -> i64:
+    return len(b)
 [out]
-def bytes_builder_basics():
+def bytes_writer_basics():
     r0, b :: librt.strings.BytesWriter
     x :: u8
     r1 :: None
     r2 :: bytes
     r3 :: None
-    r4 :: bytes
+    n :: i64
+    r4 :: None
+    r5 :: bytes
 L0:
     r0 = LibRTStrings_BytesWriter_internal()
     b = r0
     x = 1
     r1 = LibRTStrings_BytesWriter_append_internal(b, x)
     r2 = b'foo'
     r3 = LibRTStrings_BytesWriter_write_internal(b, r2)
-    r4 = LibRTStrings_BytesWriter_getvalue_internal(b)
-    return r4
+    n = 2
+    r4 = LibRTStrings_BytesWriter_truncate_internal(b, n)
+    r5 = LibRTStrings_BytesWriter_getvalue_internal(b)
+    return r5
+def bytes_writer_len(b):
+    b :: librt.strings.BytesWriter
+    r0 :: short_int
+    r1 :: i64
+L0:
+    r0 = LibRTStrings_BytesWriter_len_internal(b)
+    r1 = r0 >> 1
+    return r1