[mypyc] Speed up BytesWriter append, write and len operations

mypyc/ir/deps.py

@@ -50,3 +50,4 @@ def get_header(self) -> str:
 LIBRT_BASE64: Final = Capsule("librt.base64")
 
 BYTES_EXTRA_OPS: Final = SourceDep("bytes_extra_ops.c")
+BYTES_WRITER_EXTRA_OPS: Final = SourceDep("byteswriter_extra_ops.c")

mypyc/lib-rt/byteswriter_extra_ops.c

@@ -0,0 +1,31 @@
+// Primitives related to librt.strings.BytesWriter that get linked statically
+// with compiled modules, instead of being called via a capsule.
+
+#include "byteswriter_extra_ops.h"
+
+char CPyBytesWriter_Write(PyObject *obj, PyObject *value) {
+    BytesWriterObject *self = (BytesWriterObject *)obj;
+    const char *data;
+    Py_ssize_t size;
+    if (likely(PyBytes_Check(value))) {
+        data = PyBytes_AS_STRING(value);
+        size = PyBytes_GET_SIZE(value);
+    } else {
+        data = PyByteArray_AS_STRING(value);
+        size = PyByteArray_GET_SIZE(value);
+    }
+    // Write bytes content.
+    if (!CPyBytesWriter_EnsureSize(self, size))
+        return CPY_NONE_ERROR;
+    if (size < 8) {
+        // Loop tends to be faster for small sizes
+        char *p = self->buf + self->len;
+        for (Py_ssize_t i = 0; i < size; i++) {
+            p[i] = data[i];
+        }
+    } else {
+        memcpy(self->buf + self->len, data, size);
+    }
+    self->len += size;
+    return CPY_NONE;
+}

mypyc/lib-rt/byteswriter_extra_ops.h

@@ -0,0 +1,34 @@
+#ifndef BYTESWRITER_EXTRA_OPS_H
+#define BYTESWRITER_EXTRA_OPS_H
+
+#include "librt_strings.h"
+
+static inline CPyTagged
+CPyBytesWriter_Len(PyObject *obj) {
+    return (CPyTagged)((BytesWriterObject *)obj)->len << 1;
+}
+
+static inline bool
+CPyBytesWriter_EnsureSize(BytesWriterObject *data, Py_ssize_t n) {
+    if (likely(data->capacity - data->len >= n)) {
+        return true;
+    } else {
+        return LibRTStrings_ByteWriter_grow_buffer_internal(data, n);
+    }
+}
+
+static inline char
+CPyBytesWriter_Append(PyObject *obj, uint8_t value) {
+    BytesWriterObject *self = (BytesWriterObject *)obj;
+    // Store length in a local variable to enable additional optimizations
+    Py_ssize_t len = self->len;
+    if (!CPyBytesWriter_EnsureSize(self, 1))
+        return CPY_NONE_ERROR;
+    self->buf[len] = value;
+    self->len = len + 1;
+    return CPY_NONE;
+}
+
+char CPyBytesWriter_Write(PyObject *obj, PyObject *value);
+
+#endif

mypyc/lib-rt/librt_internal.c

@@ -27,8 +27,6 @@
 #define LONG_INT_TRAILER 15
 
 #define CPY_BOOL_ERROR 2
-#define CPY_NONE_ERROR 2
-#define CPY_NONE 1
 
 #define _CHECK_READ_BUFFER(data, err)  if (unlikely(_check_read_buffer(data) == CPY_NONE_ERROR)) \
                                            return err;

mypyc/lib-rt/librt_strings.c

@@ -7,30 +7,19 @@
 #include "librt_strings.h"
 
 #define CPY_BOOL_ERROR 2
-#define CPY_NONE_ERROR 2
-#define CPY_NONE 1
 
 //
 // BytesWriter
 //
 
-// Length of the default buffer embedded directly in a BytesWriter object
-#define WRITER_EMBEDDED_BUF_LEN 512
-
-typedef struct {
-    PyObject_HEAD
-    char *buf;  // Beginning of the buffer
-    Py_ssize_t len;  // Current length (number of bytes written)
-    Py_ssize_t capacity;  // Total capacity of the buffer
-    char data[WRITER_EMBEDDED_BUF_LEN];  // Default buffer
-} BytesWriterObject;
-
 #define _WRITE(data, type, v) \
     do { \
        *(type *)(((BytesWriterObject *)data)->buf + ((BytesWriterObject *)data)->len) = v; \
        ((BytesWriterObject *)data)->len += sizeof(type); \
     } while (0)
 
+#ifdef MYPYC_EXPERIMENTAL
+
 static PyTypeObject BytesWriterType;
 
 static bool
@@ -390,6 +379,8 @@ BytesWriter_len_internal(PyObject *self) {
     return writer->len << 1;
 }
 
+#endif
+
 static PyMethodDef librt_strings_module_methods[] = {
     {NULL, NULL, 0, NULL}
 };
@@ -426,9 +417,8 @@ librt_strings_module_exec(PyObject *m)
         (void *)BytesWriter_internal,
         (void *)BytesWriter_getvalue_internal,
         (void *)BytesWriter_append_internal,
-        (void *)BytesWriter_write_internal,
+        (void *)_grow_buffer,
         (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);

mypyc/lib-rt/librt_strings.h

@@ -15,28 +15,38 @@ import_librt_strings(void)
 // ABI version -- only an exact match is compatible. This will only be changed in
 // very exceptional cases (likely never) due to strict backward compatibility
 // requirements.
-#define LIBRT_STRINGS_ABI_VERSION 0
+#define LIBRT_STRINGS_ABI_VERSION 1
 
 // 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 1
+ #define LIBRT_STRINGS_API_VERSION 2
 
 // Number of functions in the capsule API. If you add a new function, also increase
 // LIBRT_STRINGS_API_VERSION.
-#define LIBRT_STRINGS_API_LEN 9
+#define LIBRT_STRINGS_API_LEN 8
 
 static void *LibRTStrings_API[LIBRT_STRINGS_API_LEN];
 
+// Length of the default buffer embedded directly in a BytesWriter object
+#define WRITER_EMBEDDED_BUF_LEN 512
+
+typedef struct {
+    PyObject_HEAD
+    char *buf;  // Beginning of the buffer
+    Py_ssize_t len;  // Current length (number of bytes written)
+    Py_ssize_t capacity;  // Total capacity of the buffer
+    char data[WRITER_EMBEDDED_BUF_LEN];  // Default buffer
+} BytesWriterObject;
+
 #define LibRTStrings_ABIVersion (*(int (*)(void)) LibRTStrings_API[0])
 #define LibRTStrings_APIVersion (*(int (*)(void)) LibRTStrings_API[1])
 #define LibRTStrings_BytesWriter_internal (*(PyObject* (*)(void)) LibRTStrings_API[2])
 #define LibRTStrings_BytesWriter_getvalue_internal (*(PyObject* (*)(PyObject *source)) LibRTStrings_API[3])
 #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_ByteWriter_grow_buffer_internal (*(bool (*)(BytesWriterObject *obj, Py_ssize_t size)) 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])
+#define LibRTStrings_BytesWriter_truncate_internal (*(char (*)(PyObject *self, int64_t size)) LibRTStrings_API[7])
 
 static int
 import_librt_strings(void)

mypyc/lib-rt/mypyc_util.h

@@ -142,6 +142,10 @@ typedef PyObject CPyModule;
 // Error value for floats
 #define CPY_FLOAT_ERROR -113.0
 
+// Value for 'None' primitive type
+#define CPY_NONE_ERROR 2
+#define CPY_NONE 1
+
 typedef void (*CPyVTableItem)(void);
 
 static inline CPyTagged CPyTagged_ShortFromInt(int x) {

mypyc/primitives/librt_strings_ops.py

@@ -1,6 +1,6 @@
 from typing import Final
 
-from mypyc.ir.deps import LIBRT_STRINGS
+from mypyc.ir.deps import BYTES_WRITER_EXTRA_OPS, LIBRT_STRINGS
 from mypyc.ir.ops import ERR_MAGIC, ERR_NEVER
 from mypyc.ir.rtypes import (
     KNOWN_NATIVE_TYPES,
@@ -38,20 +38,20 @@
     name="write",
     arg_types=[bytes_writer_rprimitive, bytes_rprimitive],
     return_type=none_rprimitive,
-    c_function_name="LibRTStrings_BytesWriter_write_internal",
+    c_function_name="CPyBytesWriter_Write",
     error_kind=ERR_MAGIC,
     experimental=True,
-    dependencies=[LIBRT_STRINGS],
+    dependencies=[LIBRT_STRINGS, BYTES_WRITER_EXTRA_OPS],
 )
 
 method_op(
     name="append",
     arg_types=[bytes_writer_rprimitive, uint8_rprimitive],
     return_type=none_rprimitive,
-    c_function_name="LibRTStrings_BytesWriter_append_internal",
+    c_function_name="CPyBytesWriter_Append",
     error_kind=ERR_MAGIC,
     experimental=True,
-    dependencies=[LIBRT_STRINGS],
+    dependencies=[LIBRT_STRINGS, BYTES_WRITER_EXTRA_OPS],
 )
 
 method_op(
@@ -60,14 +60,16 @@
     return_type=none_rprimitive,
     c_function_name="LibRTStrings_BytesWriter_truncate_internal",
     error_kind=ERR_MAGIC,
+    experimental=True,
+    dependencies=[LIBRT_STRINGS],
 )
 
 function_op(
     name="builtins.len",
     arg_types=[bytes_writer_rprimitive],
     return_type=short_int_rprimitive,
-    c_function_name="LibRTStrings_BytesWriter_len_internal",
+    c_function_name="CPyBytesWriter_Len",
     error_kind=ERR_NEVER,
     experimental=True,
-    dependencies=[LIBRT_STRINGS],
+    dependencies=[LIBRT_STRINGS, BYTES_WRITER_EXTRA_OPS],
 )

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

@@ -29,9 +29,9 @@ def bytes_writer_basics():
 L0:
     r0 = LibRTStrings_BytesWriter_internal()
     b = r0
-    r1 = LibRTStrings_BytesWriter_append_internal(b, 1)
+    r1 = CPyBytesWriter_Append(b, 1)
     r2 = b'foo'
-    r3 = LibRTStrings_BytesWriter_write_internal(b, r2)
+    r3 = CPyBytesWriter_Write(b, r2)
     n = 4
     r4 = n & 1
     r5 = r4 == 0
@@ -55,6 +55,6 @@ def bytes_writer_len(b):
     r0 :: short_int
     r1 :: i64
 L0:
-    r0 = LibRTStrings_BytesWriter_len_internal(b)
+    r0 = CPyBytesWriter_Len(b)
     r1 = r0 >> 1
     return r1