[mypyc] Add write_i16_le and read_i16_le to librt.strings (#20745)

These are experimental optimized primitives for reading/writing 16-bit
binary integers. I'll later add support for other integer sizes and
floats, and big endian formats.

The main optimization here is to inline the hot paths of these
functions. Also the api is optimized for performance and not really
ergonomics, though I don't think the ergonomics are particularly bad.
When reading, the caller needs to keep track of the index to the bytes
object, and negative indexes aren't supported, but the latter doesn't
seem really useful for these kinds of primitives.

In a microbenchmark, these can be dozens of times faster than using
stdlib `struct` (when compiled).

It may still be possible to further optimize performance, but these seem
already pretty good.

I used coding agent assist but through small human reviewed increments.

Author: JukkaL

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

@@ -1,6 +1,6 @@
 from typing import final
 
-from mypy_extensions import i64, i32, u8
+from mypy_extensions import i64, i32, i16, u8
 
 @final
 class BytesWriter:
@@ -19,3 +19,6 @@ class StringWriter:
     def getvalue(self) -> str: ...
     def __len__(self) -> i64: ...
     def __getitem__(self, i: i64, /) -> i32: ...
+
+def write_i16_le(b: BytesWriter, n: i16, /) -> None: ...
+def read_i16_le(b: bytes, index: i64, /) -> i16: ...

mypyc/lib-rt/byteswriter_extra_ops.c

@@ -32,4 +32,15 @@ char CPyBytesWriter_Write(PyObject *obj, PyObject *value) {
     return CPY_NONE;
 }
 
+int16_t CPyBytes_ReadError(int64_t index, Py_ssize_t size) {
+    if (index < 0) {
+        PyErr_SetString(PyExc_ValueError, "index must be non-negative");
+    } else {
+        PyErr_Format(PyExc_IndexError,
+                     "index %lld out of range for bytes of length %zd",
+                     (long long)index, size);
+    }
+    return CPY_LL_INT_ERROR;
+}
+
 #endif // MYPYC_EXPERIMENTAL

mypyc/lib-rt/byteswriter_extra_ops.h

@@ -7,6 +7,7 @@
 #include <Python.h>
 
 #include "strings/librt_strings.h"
+#include "strings/librt_strings_common.h"
 
 static inline CPyTagged
 CPyBytesWriter_Len(PyObject *obj) {
@@ -34,8 +35,20 @@ CPyBytesWriter_Append(PyObject *obj, uint8_t value) {
     return CPY_NONE;
 }
 
+static inline char
+CPyBytesWriter_WriteI16LE(PyObject *obj, int16_t value) {
+    BytesWriterObject *self = (BytesWriterObject *)obj;
+    if (!CPyBytesWriter_EnsureSize(self, 2))
+        return CPY_NONE_ERROR;
+    BytesWriter_write_i16_le_unchecked(self, value);
+    return CPY_NONE;
+}
+
 char CPyBytesWriter_Write(PyObject *obj, PyObject *value);
 
+// Helper function for bytes read error handling (negative index or out of range)
+int16_t CPyBytes_ReadError(int64_t index, Py_ssize_t size);
+
 // If index is negative, convert to non-negative index (no range checking)
 static inline int64_t CPyBytesWriter_AdjustIndex(PyObject *obj, int64_t index) {
     if (index < 0) {
@@ -56,6 +69,18 @@ static inline void CPyBytesWriter_SetItem(PyObject *obj, int64_t index, uint8_t
     (((BytesWriterObject *)obj)->buf)[index] = x;
 }
 
+static inline int16_t
+CPyBytes_ReadI16LE(PyObject *bytes_obj, int64_t index) {
+    // bytes_obj type is enforced by mypyc
+    Py_ssize_t size = PyBytes_GET_SIZE(bytes_obj);
+    if (unlikely(index < 0 || index > size - 2)) {
+        CPyBytes_ReadError(index, size);
+        return CPY_LL_INT_ERROR;
+    }
+    const unsigned char *data = (const unsigned char *)PyBytes_AS_STRING(bytes_obj);
+    return read_i16_le_unchecked(data + index);
+}
+
 #endif // MYPYC_EXPERIMENTAL
 
 #endif

mypyc/lib-rt/strings/librt_strings.c

@@ -831,9 +831,76 @@ StringWriter_len_internal(PyObject *self) {
 
 // End of StringWriter
 
+static PyObject*
+write_i16_le(PyObject *module, PyObject *const *args, size_t nargs) {
+    if (unlikely(nargs != 2)) {
+        PyErr_Format(PyExc_TypeError,
+                     "write_i16_le() takes exactly 2 arguments (%zu given)", nargs);
+        return NULL;
+    }
+    PyObject *writer = args[0];
+    if (!check_bytes_writer(writer)) {
+        return NULL;
+    }
+    PyObject *value = args[1];
+    int16_t unboxed = CPyLong_AsInt16(value);
+    if (unlikely(unboxed == CPY_LL_INT_ERROR && PyErr_Occurred())) {
+        // Error already set by CPyLong_AsInt16 (ValueError for overflow, TypeError for wrong type)
+        return NULL;
+    }
+    BytesWriterObject *bw = (BytesWriterObject *)writer;
+    if (unlikely(!ensure_bytes_writer_size(bw, 2))) {
+        return NULL;
+    }
+    BytesWriter_write_i16_le_unchecked(bw, unboxed);
+    Py_INCREF(Py_None);
+    return Py_None;
+}
+
+static PyObject*
+read_i16_le(PyObject *module, PyObject *const *args, size_t nargs) {
+    if (unlikely(nargs != 2)) {
+        PyErr_Format(PyExc_TypeError,
+                     "read_i16_le() takes exactly 2 arguments (%zu given)", nargs);
+        return NULL;
+    }
+    PyObject *bytes_obj = args[0];
+    if (unlikely(!PyBytes_Check(bytes_obj))) {
+        PyErr_SetString(PyExc_TypeError, "read_i16_le() argument 1 must be bytes");
+        return NULL;
+    }
+    PyObject *index_obj = args[1];
+    int64_t index = CPyLong_AsInt64(index_obj);
+    if (unlikely(index == CPY_LL_INT_ERROR && PyErr_Occurred())) {
+        return NULL;
+    }
+    if (unlikely(index < 0)) {
+        PyErr_SetString(PyExc_ValueError, "index must be non-negative");
+        return NULL;
+    }
+    Py_ssize_t size = PyBytes_GET_SIZE(bytes_obj);
+    if (unlikely(index > size - 2)) {
+        PyErr_Format(PyExc_IndexError,
+                     "index %lld out of range for bytes of length %zd",
+                     (long long)index, size);
+        return NULL;
+    }
+    const unsigned char *data = (const unsigned char *)PyBytes_AS_STRING(bytes_obj);
+    int16_t value = read_i16_le_unchecked(data + index);
+    return PyLong_FromLong(value);
+}
+
 #endif
 
 static PyMethodDef librt_strings_module_methods[] = {
+#ifdef MYPYC_EXPERIMENTAL
+    {"write_i16_le", (PyCFunction) write_i16_le, METH_FASTCALL,
+     PyDoc_STR("Write a 16-bit signed integer to BytesWriter in little-endian format")
+    },
+    {"read_i16_le", (PyCFunction) read_i16_le, METH_FASTCALL,
+     PyDoc_STR("Read a 16-bit signed integer from bytes in little-endian format")
+    },
+#endif
     {NULL, NULL, 0, NULL}
 };
 

mypyc/lib-rt/strings/librt_strings.h

@@ -13,6 +13,7 @@ import_librt_strings(void)
 #else  // MYPYC_EXPERIMENTAL
 
 #include <Python.h>
+#include "librt_strings_common.h"
 
 // ABI version -- only an exact match is compatible. This will only be changed in
 // very exceptional cases (likely never) due to strict backward compatibility
@@ -30,17 +31,6 @@ import_librt_strings(void)
 
 static void *LibRTStrings_API[LIBRT_STRINGS_API_LEN];
 
-// Length of the default buffer embedded directly in a BytesWriter object
-#define WRITER_EMBEDDED_BUF_LEN 256
-
-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;
-
 typedef struct {
     PyObject_HEAD
     char *buf;  // Beginning of the buffer

mypyc/lib-rt/strings/librt_strings_common.h

@@ -0,0 +1,46 @@
+#ifndef LIBRT_STRINGS_COMMON_H
+#define LIBRT_STRINGS_COMMON_H
+
+#include <Python.h>
+#include <stdint.h>
+#include <string.h>
+
+// Length of the default buffer embedded directly in a BytesWriter object
+#define WRITER_EMBEDDED_BUF_LEN 256
+
+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;
+
+// Write a 16-bit signed integer in little-endian format to BytesWriter.
+// NOTE: This does NOT check buffer capacity - caller must ensure space is available.
+static inline void
+BytesWriter_write_i16_le_unchecked(BytesWriterObject *self, int16_t value) {
+    // Store len in local to help optimizer reduce struct member accesses
+    Py_ssize_t len = self->len;
+    unsigned char *p = (unsigned char *)(self->buf + len);
+    uint16_t uval = (uint16_t)value;
+
+    // Write in little-endian format
+    // Modern compilers optimize this pattern well, often to a single store on LE systems
+    p[0] = (unsigned char)uval;
+    p[1] = (unsigned char)(uval >> 8);
+
+    self->len = len + 2;
+}
+
+// Read a 16-bit signed integer in little-endian format from bytes.
+// NOTE: This does NOT check bounds - caller must ensure valid index.
+static inline int16_t
+read_i16_le_unchecked(const unsigned char *data) {
+    // Read in little-endian format
+    // Modern compilers optimize this pattern well, often to a single load on LE systems
+    uint16_t uval = (uint16_t)data[0] | ((uint16_t)data[1] << 8);
+    return (int16_t)uval;
+}
+
+#endif  // LIBRT_STRINGS_COMMON_H

mypyc/primitives/librt_strings_ops.py

@@ -5,6 +5,7 @@
     bytearray_rprimitive,
     bytes_rprimitive,
     bytes_writer_rprimitive,
+    int16_rprimitive,
     int32_rprimitive,
     int64_rprimitive,
     none_rprimitive,
@@ -76,6 +77,26 @@
     dependencies=[LIBRT_STRINGS],
 )
 
+function_op(
+    name="librt.strings.write_i16_le",
+    arg_types=[bytes_writer_rprimitive, int16_rprimitive],
+    return_type=none_rprimitive,
+    c_function_name="CPyBytesWriter_WriteI16LE",
+    error_kind=ERR_MAGIC,
+    experimental=True,
+    dependencies=[LIBRT_STRINGS, BYTES_WRITER_EXTRA_OPS],
+)
+
+function_op(
+    name="librt.strings.read_i16_le",
+    arg_types=[bytes_rprimitive, int64_rprimitive],
+    return_type=int16_rprimitive,
+    c_function_name="CPyBytes_ReadI16LE",
+    error_kind=ERR_MAGIC,
+    experimental=True,
+    dependencies=[LIBRT_STRINGS, BYTES_WRITER_EXTRA_OPS],
+)
+
 function_op(
     name="builtins.len",
     arg_types=[bytes_writer_rprimitive],

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

@@ -94,6 +94,52 @@ L2:
     CPyBytesWriter_SetItem(b, r0, x)
     return 1
 
+[case testLibrtStrings_write_i16_le_experimental_64bit]
+from librt.strings import BytesWriter, write_i16_le
+from mypy_extensions import i16
+
+def test_write_i16_le(b: BytesWriter, n: i16) -> None:
+    write_i16_le(b, n)
+def test_write_i16_le_literal(b: BytesWriter) -> None:
+    write_i16_le(b, 1234)
+[out]
+def test_write_i16_le(b, n):
+    b :: librt.strings.BytesWriter
+    n :: i16
+    r0 :: None
+L0:
+    r0 = CPyBytesWriter_WriteI16LE(b, n)
+    return 1
+def test_write_i16_le_literal(b):
+    b :: librt.strings.BytesWriter
+    r0 :: None
+L0:
+    r0 = CPyBytesWriter_WriteI16LE(b, 1234)
+    return 1
+
+[case testLibrtStrings_read_i16_le_experimental_64bit]
+from librt.strings import read_i16_le
+from mypy_extensions import i16, i64
+
+def test_read_i16_le(b: bytes, i: i64) -> i16:
+    return read_i16_le(b, i)
+def test_read_i16_le_literal(b: bytes) -> i16:
+    return read_i16_le(b, 0)
+[out]
+def test_read_i16_le(b, i):
+    b :: bytes
+    i :: i64
+    r0 :: i16
+L0:
+    r0 = CPyBytes_ReadI16LE(b, i)
+    return r0
+def test_read_i16_le_literal(b):
+    b :: bytes
+    r0 :: i16
+L0:
+    r0 = CPyBytes_ReadI16LE(b, 0)
+    return r0
+
 [case testLibrtStrings_StringWriter_experimental_64bit]
 from librt.strings import StringWriter
 from mypy_extensions import i32, i64