Sync from mypy and bump version to 0.3.0

Also update smoke test. Includes `librt.internal.{write_bytes,read_bytes}`.

Author: JukkaL

lib-rt/CPy.h

@@ -148,6 +148,7 @@ CPyTagged CPyTagged_Remainder_(CPyTagged left, CPyTagged right);
 CPyTagged CPyTagged_BitwiseLongOp_(CPyTagged a, CPyTagged b, char op);
 CPyTagged CPyTagged_Rshift_(CPyTagged left, CPyTagged right);
 CPyTagged CPyTagged_Lshift_(CPyTagged left, CPyTagged right);
+CPyTagged CPyTagged_BitLength(CPyTagged self);
 
 PyObject *CPyTagged_Str(CPyTagged n);
 CPyTagged CPyTagged_FromFloat(double f);
@@ -734,6 +735,7 @@ static inline char CPyDict_CheckSize(PyObject *dict, Py_ssize_t size) {
 #define BOTHSTRIP  2
 
 char CPyStr_Equal(PyObject *str1, PyObject *str2);
+char CPyStr_EqualLiteral(PyObject *str, PyObject *literal_str, Py_ssize_t literal_length);
 PyObject *CPyStr_Build(Py_ssize_t len, ...);
 PyObject *CPyStr_GetItem(PyObject *str, CPyTagged index);
 PyObject *CPyStr_GetItemUnsafe(PyObject *str, Py_ssize_t index);

lib-rt/int_ops.c

@@ -5,6 +5,10 @@
 #include <Python.h>
 #include "CPy.h"
 
+#ifdef _MSC_VER
+#include <intrin.h>
+#endif
+
 #ifndef _WIN32
 // On 64-bit Linux and macOS, ssize_t and long are both 64 bits, and
 // PyLong_FromLong is faster than PyLong_FromSsize_t, so use the faster one
@@ -15,6 +19,17 @@
 #define CPyLong_FromSsize_t PyLong_FromSsize_t
 #endif
 
+#if defined(__GNUC__) || defined(__clang__)
+#  if defined(__x86_64__) || defined(_M_X64) || defined(__aarch64__) || (defined(__SIZEOF_POINTER__) && __SIZEOF_POINTER__ == 8)
+#    define CPY_CLZ(x) __builtin_clzll((unsigned long long)(x))
+#    define CPY_BITS 64
+#  else
+#    define CPY_CLZ(x) __builtin_clz((unsigned int)(x))
+#    define CPY_BITS 32
+#  endif
+#endif
+
+
 CPyTagged CPyTagged_FromSsize_t(Py_ssize_t value) {
     // We use a Python object if the value shifted left by 1 is too
     // large for Py_ssize_t
@@ -581,3 +596,52 @@ double CPyTagged_TrueDivide(CPyTagged x, CPyTagged y) {
     }
     return 1.0;
 }
+
+// int.bit_length()
+CPyTagged CPyTagged_BitLength(CPyTagged self) {
+    // Handle zero
+    if (self == 0) {
+        return 0;
+    }
+
+    // Fast path for small (tagged) ints
+    if (CPyTagged_CheckShort(self)) {
+        Py_ssize_t val = CPyTagged_ShortAsSsize_t(self);
+        Py_ssize_t absval = val < 0 ? -val : val;
+        int bits = 0;
+        if (absval) {
+#if defined(_MSC_VER)
+    #if defined(_WIN64)
+            unsigned long idx;
+            if (_BitScanReverse64(&idx, (unsigned __int64)absval)) {
+                bits = (int)(idx + 1);
+            }
+    #else
+            unsigned long idx;
+            if (_BitScanReverse(&idx, (unsigned long)absval)) {
+                bits = (int)(idx + 1);
+            }
+    #endif
+#elif defined(__GNUC__) || defined(__clang__)
+            bits = (int)(CPY_BITS - CPY_CLZ(absval));
+#else
+            // Fallback to loop if no builtin
+            while (absval) {
+                absval >>= 1;
+                bits++;
+            }
+#endif
+        }
+        return bits << 1;
+    }
+
+    // Slow path for big ints
+    PyObject *pyint = CPyTagged_AsObject(self);
+    int bits = _PyLong_NumBits(pyint);
+    Py_DECREF(pyint);
+    if (bits < 0) {
+        // _PyLong_NumBits sets an error on failure
+        return CPY_INT_TAG;
+    }
+    return bits << 1;
+}

lib-rt/librt_internal.c

@@ -346,6 +346,100 @@ write_str(PyObject *self, PyObject *const *args, size_t nargs, PyObject *kwnames
     return Py_None;
 }
 
+/*
+bytes format: size followed by bytes
+    short bytes (len <= 127): single byte for size as `(uint8_t)size << 1`
+    long bytes: \x01 followed by size as Py_ssize_t
+*/
+
+static PyObject*
+read_bytes_internal(PyObject *data) {
+    _CHECK_BUFFER(data, NULL)
+
+    // Read length.
+    Py_ssize_t size;
+    _CHECK_READ(data, 1, NULL)
+    uint8_t first = _READ(data, uint8_t)
+    if (likely(first != LONG_STR_TAG)) {
+        // Common case: short bytes (len <= 127).
+        size = (Py_ssize_t)(first >> 1);
+    } else {
+        _CHECK_READ(data, sizeof(CPyTagged), NULL)
+        size = _READ(data, Py_ssize_t)
+    }
+    // Read bytes content.
+    char *buf = ((BufferObject *)data)->buf;
+    _CHECK_READ(data, size, NULL)
+    PyObject *res = PyBytes_FromStringAndSize(
+        buf + ((BufferObject *)data)->pos, (Py_ssize_t)size
+    );
+    if (unlikely(res == NULL))
+        return NULL;
+    ((BufferObject *)data)->pos += size;
+    return res;
+}
+
+static PyObject*
+read_bytes(PyObject *self, PyObject *const *args, size_t nargs, PyObject *kwnames) {
+    static const char * const kwlist[] = {"data", 0};
+    static CPyArg_Parser parser = {"O:read_bytes", kwlist, 0};
+    PyObject *data;
+    if (unlikely(!CPyArg_ParseStackAndKeywordsOneArg(args, nargs, kwnames, &parser, &data))) {
+        return NULL;
+    }
+    return read_bytes_internal(data);
+}
+
+static char
+write_bytes_internal(PyObject *data, PyObject *value) {
+    _CHECK_BUFFER(data, CPY_NONE_ERROR)
+
+    const char *chunk = PyBytes_AsString(value);
+    if (unlikely(chunk == NULL))
+        return CPY_NONE_ERROR;
+    Py_ssize_t size = PyBytes_GET_SIZE(value);
+
+    Py_ssize_t need;
+    // Write length.
+    if (likely(size <= MAX_SHORT_LEN)) {
+        // Common case: short bytes (len <= 127) store as single byte.
+        need = size + 1;
+        _CHECK_SIZE(data, need)
+        _WRITE(data, uint8_t, (uint8_t)size << 1)
+    } else {
+        need = size + sizeof(Py_ssize_t) + 1;
+        _CHECK_SIZE(data, need)
+        _WRITE(data, uint8_t, LONG_STR_TAG)
+        _WRITE(data, Py_ssize_t, size)
+    }
+    // Write bytes content.
+    char *buf = ((BufferObject *)data)->buf;
+    memcpy(buf + ((BufferObject *)data)->pos, chunk, size);
+    ((BufferObject *)data)->pos += size;
+    ((BufferObject *)data)->end += need;
+    return CPY_NONE;
+}
+
+static PyObject*
+write_bytes(PyObject *self, PyObject *const *args, size_t nargs, PyObject *kwnames) {
+    static const char * const kwlist[] = {"data", "value", 0};
+    static CPyArg_Parser parser = {"OO:write_bytes", kwlist, 0};
+    PyObject *data;
+    PyObject *value;
+    if (unlikely(!CPyArg_ParseStackAndKeywordsSimple(args, nargs, kwnames, &parser, &data, &value))) {
+        return NULL;
+    }
+    if (unlikely(!PyBytes_Check(value))) {
+        PyErr_SetString(PyExc_TypeError, "value must be a bytes object");
+        return NULL;
+    }
+    if (unlikely(write_bytes_internal(data, value) == CPY_NONE_ERROR)) {
+        return NULL;
+    }
+    Py_INCREF(Py_None);
+    return Py_None;
+}
+
 /*
 float format:
     stored as a C double
@@ -565,6 +659,8 @@ static PyMethodDef librt_internal_module_methods[] = {
     {"read_bool", (PyCFunction)read_bool, METH_FASTCALL | METH_KEYWORDS, PyDoc_STR("read a bool")},
     {"write_str", (PyCFunction)write_str, METH_FASTCALL | METH_KEYWORDS, PyDoc_STR("write a string")},
     {"read_str", (PyCFunction)read_str, METH_FASTCALL | METH_KEYWORDS, PyDoc_STR("read a string")},
+    {"write_bytes", (PyCFunction)write_bytes, METH_FASTCALL | METH_KEYWORDS, PyDoc_STR("write bytes")},
+    {"read_bytes", (PyCFunction)read_bytes, METH_FASTCALL | METH_KEYWORDS, PyDoc_STR("read bytes")},
     {"write_float", (PyCFunction)write_float, METH_FASTCALL | METH_KEYWORDS, PyDoc_STR("write a float")},
     {"read_float", (PyCFunction)read_float, METH_FASTCALL | METH_KEYWORDS, PyDoc_STR("read a float")},
     {"write_int", (PyCFunction)write_int, METH_FASTCALL | METH_KEYWORDS, PyDoc_STR("write an int")},
@@ -590,7 +686,7 @@ librt_internal_module_exec(PyObject *m)
     }
 
     // Export mypy internal C API, be careful with the order!
-    static void *NativeInternal_API[14] = {
+    static void *NativeInternal_API[16] = {
         (void *)Buffer_internal,
         (void *)Buffer_internal_empty,
         (void *)Buffer_getvalue_internal,
@@ -605,6 +701,8 @@ librt_internal_module_exec(PyObject *m)
         (void *)write_tag_internal,
         (void *)read_tag_internal,
         (void *)NativeInternal_ABI_Version,
+        (void *)write_bytes_internal,
+        (void *)read_bytes_internal,
     };
     PyObject *c_api_object = PyCapsule_New((void *)NativeInternal_API, "librt.internal._C_API", NULL);
     if (PyModule_Add(m, "_C_API", c_api_object) < 0) {

lib-rt/librt_internal.h

@@ -19,6 +19,8 @@ static CPyTagged read_int_internal(PyObject *data);
 static char write_tag_internal(PyObject *data, uint8_t value);
 static uint8_t read_tag_internal(PyObject *data);
 static int NativeInternal_ABI_Version(void);
+static char write_bytes_internal(PyObject *data, PyObject *value);
+static PyObject *read_bytes_internal(PyObject *data);
 
 #else
 
@@ -38,6 +40,8 @@ static void **NativeInternal_API;
 #define write_tag_internal (*(char (*)(PyObject *source, uint8_t value)) NativeInternal_API[11])
 #define read_tag_internal (*(uint8_t (*)(PyObject *source)) NativeInternal_API[12])
 #define NativeInternal_ABI_Version (*(int (*)(void)) NativeInternal_API[13])
+#define write_bytes_internal (*(char (*)(PyObject *source, PyObject *value)) NativeInternal_API[14])
+#define read_bytes_internal (*(PyObject* (*)(PyObject *source)) NativeInternal_API[15])
 
 static int
 import_librt_internal(void)

lib-rt/setup.py

@@ -1,4 +1,4 @@
-"""Build script for mypyc C runtime library unit tests.
+"""Build script for mypyc C runtime library and C API unit tests.
 
 The tests are written in C++ and use the Google Test framework.
 """

lib-rt/str_ops.c

@@ -64,20 +64,33 @@ make_bloom_mask(int kind, const void* ptr, Py_ssize_t len)
 #undef BLOOM_UPDATE
 }
 
-// Adapted from CPython 3.13.1 (_PyUnicode_Equal)
-char CPyStr_Equal(PyObject *str1, PyObject *str2) {
-    if (str1 == str2) {
-        return 1;
-    }
-    Py_ssize_t len = PyUnicode_GET_LENGTH(str1);
-    if (PyUnicode_GET_LENGTH(str2) != len)
+static inline char _CPyStr_Equal_NoIdentCheck(PyObject *str1, PyObject *str2, Py_ssize_t str2_length) {
+    // This helper function only exists to deduplicate code in CPyStr_Equal and CPyStr_EqualLiteral
+    Py_ssize_t str1_length = PyUnicode_GET_LENGTH(str1);
+    if (str1_length != str2_length)
         return 0;
     int kind = PyUnicode_KIND(str1);
     if (PyUnicode_KIND(str2) != kind)
         return 0;
     const void *data1 = PyUnicode_DATA(str1);
     const void *data2 = PyUnicode_DATA(str2);
-    return memcmp(data1, data2, len * kind) == 0;
+    return memcmp(data1, data2, str1_length * kind) == 0;
+}
+
+// Adapted from CPython 3.13.1 (_PyUnicode_Equal)
+char CPyStr_Equal(PyObject *str1, PyObject *str2) {
+    if (str1 == str2) {
+        return 1;
+    }
+    Py_ssize_t str2_length = PyUnicode_GET_LENGTH(str2);
+    return _CPyStr_Equal_NoIdentCheck(str1, str2, str2_length);
+}
+
+char CPyStr_EqualLiteral(PyObject *str, PyObject *literal_str, Py_ssize_t literal_length) {
+    if (str == literal_str) {
+        return 1;
+    }
+    return _CPyStr_Equal_NoIdentCheck(str, literal_str, literal_length);
 }
 
 PyObject *CPyStr_GetItem(PyObject *str, CPyTagged index) {

pyproject.toml

@@ -19,7 +19,7 @@ authors = [
     {name = "Jukka Lehtosalo", email = "jukka.lehtosalo@iki.fi"},
     {name = "Ivan Levkivskyi", email = "levkivskyi@gmail.com"},
 ]
-version = "0.2.3"
+version = "0.3.0"
 license = {text = "MIT"}
 classifiers = [
   "Development Status :: 3 - Alpha",

smoke_tests.py

@@ -8,6 +8,7 @@
     write_float, read_float,
     write_int, read_int,
     write_tag, read_tag,
+    write_bytes, read_bytes,
 )
 
 Tag = u8
@@ -43,6 +44,8 @@ def test_buffer_roundtrip() -> None:
     write_int(b, 2 ** 85)
     write_int(b, 1234512344)
     write_int(b, 1234512345)
+    write_bytes(b, b"foobar")
+    write_bytes(b, b"abc" * 1000)
 
     b = Buffer(b.getvalue())
     assert read_str(b) == "foo"
@@ -58,6 +61,8 @@ def test_buffer_roundtrip() -> None:
     assert read_int(b) == 2 ** 85
     assert read_int(b) == 1234512344
     assert read_int(b) == 1234512345
+    assert read_bytes(b) == b"foobar"
+    assert read_bytes(b) == b"abc" * 1000
 
 
 def test_buffer_int_size() -> None: