diff --git a/mypyc/ir/deps.py b/mypyc/ir/deps.py index b7747646628b..9aee5d1b33e4 100644 --- a/mypyc/ir/deps.py +++ b/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") diff --git a/mypyc/lib-rt/byteswriter_extra_ops.c b/mypyc/lib-rt/byteswriter_extra_ops.c new file mode 100644 index 000000000000..03827bf54fea --- /dev/null +++ b/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; +} diff --git a/mypyc/lib-rt/byteswriter_extra_ops.h b/mypyc/lib-rt/byteswriter_extra_ops.h new file mode 100644 index 000000000000..09cd018d4e5f --- /dev/null +++ b/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 diff --git a/mypyc/lib-rt/librt_internal.c b/mypyc/lib-rt/librt_internal.c index fe18c541c11f..9c57d8bef72c 100644 --- a/mypyc/lib-rt/librt_internal.c +++ b/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; diff --git a/mypyc/lib-rt/librt_strings.c b/mypyc/lib-rt/librt_strings.c index dcad4137e7f1..d97d0f571f5b 100644 --- a/mypyc/lib-rt/librt_strings.c +++ b/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); diff --git a/mypyc/lib-rt/librt_strings.h b/mypyc/lib-rt/librt_strings.h index 75eae8451e55..e619435f5292 100644 --- a/mypyc/lib-rt/librt_strings.h +++ b/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) diff --git a/mypyc/lib-rt/mypyc_util.h b/mypyc/lib-rt/mypyc_util.h index 50a806c91a8a..6715d67d9657 100644 --- a/mypyc/lib-rt/mypyc_util.h +++ b/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) { diff --git a/mypyc/primitives/librt_strings_ops.py b/mypyc/primitives/librt_strings_ops.py index 786784d15f60..1120254e24ae 100644 --- a/mypyc/primitives/librt_strings_ops.py +++ b/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], ) diff --git a/mypyc/test-data/irbuild-librt-strings.test b/mypyc/test-data/irbuild-librt-strings.test index 21aba382cdd8..b61649844de2 100644 --- a/mypyc/test-data/irbuild-librt-strings.test +++ b/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