[mypyc] Add i32 and i64 read/write functions to librt.strings (#20757)

Also fix a bug related to overlapping error values. Refactor code and
tests. Use memcpy() for unaligned access instead of accessing invidual
bytes, since this apparently works efficiently with more C compilers.

I'll add big endian variants in a follow-up PR.

I used coding agent assist, but used it to make small incremental
changes.

Follow-up to #20745.

Author: JukkaL

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

@@ -22,3 +22,7 @@ class StringWriter:
 
 def write_i16_le(b: BytesWriter, n: i16, /) -> None: ...
 def read_i16_le(b: bytes, index: i64, /) -> i16: ...
+def write_i32_le(b: BytesWriter, n: i32, /) -> None: ...
+def read_i32_le(b: bytes, index: i64, /) -> i32: ...
+def write_i64_le(b: BytesWriter, n: i64, /) -> None: ...
+def read_i64_le(b: bytes, index: i64, /) -> i64: ...

mypyc/lib-rt/byteswriter_extra_ops.c

@@ -32,15 +32,14 @@ char CPyBytesWriter_Write(PyObject *obj, PyObject *value) {
     return CPY_NONE;
 }
 
-int16_t CPyBytes_ReadError(int64_t index, Py_ssize_t size) {
+void 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

@@ -9,6 +9,8 @@
 #include "strings/librt_strings.h"
 #include "strings/librt_strings_common.h"
 
+// BytesWriter: Length and capacity
+
 static inline CPyTagged
 CPyBytesWriter_Len(PyObject *obj) {
     return (CPyTagged)((BytesWriterObject *)obj)->len << 1;
@@ -23,6 +25,8 @@ CPyBytesWriter_EnsureSize(BytesWriterObject *data, Py_ssize_t n) {
     }
 }
 
+// BytesWriter: Basic write operations
+
 static inline char
 CPyBytesWriter_Append(PyObject *obj, uint8_t value) {
     BytesWriterObject *self = (BytesWriterObject *)obj;
@@ -35,19 +39,9 @@ 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);
+// BytesWriter: Indexing operations
 
 // If index is negative, convert to non-negative index (no range checking)
 static inline int64_t CPyBytesWriter_AdjustIndex(PyObject *obj, int64_t index) {
@@ -69,6 +63,40 @@ static inline void CPyBytesWriter_SetItem(PyObject *obj, int64_t index, uint8_t
     (((BytesWriterObject *)obj)->buf)[index] = x;
 }
 
+// BytesWriter: Write integer operations (little-endian)
+
+static inline char
+CPyBytesWriter_WriteI16LE(PyObject *obj, int16_t value) {
+    BytesWriterObject *self = (BytesWriterObject *)obj;
+    if (!CPyBytesWriter_EnsureSize(self, 2))
+        return CPY_NONE_ERROR;
+    BytesWriter_WriteI16LEUnsafe(self, value);
+    return CPY_NONE;
+}
+
+static inline char
+CPyBytesWriter_WriteI32LE(PyObject *obj, int32_t value) {
+    BytesWriterObject *self = (BytesWriterObject *)obj;
+    if (!CPyBytesWriter_EnsureSize(self, 4))
+        return CPY_NONE_ERROR;
+    BytesWriter_WriteI32LEUnsafe(self, value);
+    return CPY_NONE;
+}
+
+static inline char
+CPyBytesWriter_WriteI64LE(PyObject *obj, int64_t value) {
+    BytesWriterObject *self = (BytesWriterObject *)obj;
+    if (!CPyBytesWriter_EnsureSize(self, 8))
+        return CPY_NONE_ERROR;
+    BytesWriter_WriteI64LEUnsafe(self, value);
+    return CPY_NONE;
+}
+
+// Bytes: Read integer operations (little-endian)
+
+// Helper function for bytes read error handling (negative index or out of range)
+void CPyBytes_ReadError(int64_t index, Py_ssize_t size);
+
 static inline int16_t
 CPyBytes_ReadI16LE(PyObject *bytes_obj, int64_t index) {
     // bytes_obj type is enforced by mypyc
@@ -78,7 +106,31 @@ CPyBytes_ReadI16LE(PyObject *bytes_obj, int64_t index) {
         return CPY_LL_INT_ERROR;
     }
     const unsigned char *data = (const unsigned char *)PyBytes_AS_STRING(bytes_obj);
-    return read_i16_le_unchecked(data + index);
+    return CPyBytes_ReadI16LEUnsafe(data + index);
+}
+
+static inline int32_t
+CPyBytes_ReadI32LE(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 - 4)) {
+        CPyBytes_ReadError(index, size);
+        return CPY_LL_INT_ERROR;
+    }
+    const unsigned char *data = (const unsigned char *)PyBytes_AS_STRING(bytes_obj);
+    return CPyBytes_ReadI32LEUnsafe(data + index);
+}
+
+static inline int64_t
+CPyBytes_ReadI64LE(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 - 8)) {
+        CPyBytes_ReadError(index, size);
+        return CPY_LL_INT_ERROR;
+    }
+    const unsigned char *data = (const unsigned char *)PyBytes_AS_STRING(bytes_obj);
+    return CPyBytes_ReadI64LEUnsafe(data + index);
 }
 
 #endif // MYPYC_EXPERIMENTAL

mypyc/lib-rt/strings/librt_strings.c

@@ -831,46 +831,37 @@ StringWriter_len_internal(PyObject *self) {
 
 // End of StringWriter
 
-static PyObject*
-write_i16_le(PyObject *module, PyObject *const *args, size_t nargs) {
+// Helper for write_i*_le/be functions - validates args and returns BytesWriter
+static inline BytesWriterObject *
+parse_write_int_args(PyObject *const *args, size_t nargs, const char *func_name) {
     if (unlikely(nargs != 2)) {
         PyErr_Format(PyExc_TypeError,
-                     "write_i16_le() takes exactly 2 arguments (%zu given)", nargs);
+                     "%s() takes exactly 2 arguments (%zu given)", func_name, 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;
+    return (BytesWriterObject *)writer;
 }
 
-static PyObject*
-read_i16_le(PyObject *module, PyObject *const *args, size_t nargs) {
+// Helper for read_i*_le/be functions - validates args and returns data pointer
+// Returns NULL on error, sets *out_index to the validated index on success
+static inline const unsigned char *
+parse_read_int_args(PyObject *const *args, size_t nargs, const char *func_name,
+                    Py_ssize_t num_bytes, int64_t *out_index) {
     if (unlikely(nargs != 2)) {
         PyErr_Format(PyExc_TypeError,
-                     "read_i16_le() takes exactly 2 arguments (%zu given)", nargs);
+                     "%s() takes exactly 2 arguments (%zu given)", func_name, 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");
+        PyErr_Format(PyExc_TypeError, "%s() argument 1 must be bytes", func_name);
         return NULL;
     }
-    PyObject *index_obj = args[1];
-    int64_t index = CPyLong_AsInt64(index_obj);
+    int64_t index = CPyLong_AsInt64(args[1]);
     if (unlikely(index == CPY_LL_INT_ERROR && PyErr_Occurred())) {
         return NULL;
     }
@@ -879,15 +870,83 @@ read_i16_le(PyObject *module, PyObject *const *args, size_t nargs) {
         return NULL;
     }
     Py_ssize_t size = PyBytes_GET_SIZE(bytes_obj);
-    if (unlikely(index > size - 2)) {
+    if (unlikely(index > size - num_bytes)) {
         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);
+    *out_index = index;
+    return (const unsigned char *)PyBytes_AS_STRING(bytes_obj);
+}
+
+static PyObject*
+write_i16_le(PyObject *module, PyObject *const *args, size_t nargs) {
+    BytesWriterObject *bw = parse_write_int_args(args, nargs, "write_i16_le");
+    if (bw == NULL)
+        return NULL;
+    int16_t unboxed = CPyLong_AsInt16(args[1]);
+    if (unlikely(unboxed == CPY_LL_INT_ERROR && PyErr_Occurred()))
+        return NULL;
+    if (unlikely(!ensure_bytes_writer_size(bw, 2)))
+        return NULL;
+    BytesWriter_WriteI16LEUnsafe(bw, unboxed);
+    Py_RETURN_NONE;
+}
+
+static PyObject*
+read_i16_le(PyObject *module, PyObject *const *args, size_t nargs) {
+    int64_t index;
+    const unsigned char *data = parse_read_int_args(args, nargs, "read_i16_le", 2, &index);
+    if (data == NULL)
+        return NULL;
+    return PyLong_FromLong(CPyBytes_ReadI16LEUnsafe(data + index));
+}
+
+static PyObject*
+write_i32_le(PyObject *module, PyObject *const *args, size_t nargs) {
+    BytesWriterObject *bw = parse_write_int_args(args, nargs, "write_i32_le");
+    if (bw == NULL)
+        return NULL;
+    int32_t unboxed = CPyLong_AsInt32(args[1]);
+    if (unlikely(unboxed == CPY_LL_INT_ERROR && PyErr_Occurred()))
+        return NULL;
+    if (unlikely(!ensure_bytes_writer_size(bw, 4)))
+        return NULL;
+    BytesWriter_WriteI32LEUnsafe(bw, unboxed);
+    Py_RETURN_NONE;
+}
+
+static PyObject*
+read_i32_le(PyObject *module, PyObject *const *args, size_t nargs) {
+    int64_t index;
+    const unsigned char *data = parse_read_int_args(args, nargs, "read_i32_le", 4, &index);
+    if (data == NULL)
+        return NULL;
+    return PyLong_FromLong(CPyBytes_ReadI32LEUnsafe(data + index));
+}
+
+static PyObject*
+write_i64_le(PyObject *module, PyObject *const *args, size_t nargs) {
+    BytesWriterObject *bw = parse_write_int_args(args, nargs, "write_i64_le");
+    if (bw == NULL)
+        return NULL;
+    int64_t unboxed = CPyLong_AsInt64(args[1]);
+    if (unlikely(unboxed == CPY_LL_INT_ERROR && PyErr_Occurred()))
+        return NULL;
+    if (unlikely(!ensure_bytes_writer_size(bw, 8)))
+        return NULL;
+    BytesWriter_WriteI64LEUnsafe(bw, unboxed);
+    Py_RETURN_NONE;
+}
+
+static PyObject*
+read_i64_le(PyObject *module, PyObject *const *args, size_t nargs) {
+    int64_t index;
+    const unsigned char *data = parse_read_int_args(args, nargs, "read_i64_le", 8, &index);
+    if (data == NULL)
+        return NULL;
+    return PyLong_FromLongLong(CPyBytes_ReadI64LEUnsafe(data + index));
 }
 
 #endif
@@ -900,6 +959,18 @@ static PyMethodDef librt_strings_module_methods[] = {
     {"read_i16_le", (PyCFunction) read_i16_le, METH_FASTCALL,
      PyDoc_STR("Read a 16-bit signed integer from bytes in little-endian format")
     },
+    {"write_i32_le", (PyCFunction) write_i32_le, METH_FASTCALL,
+     PyDoc_STR("Write a 32-bit signed integer to BytesWriter in little-endian format")
+    },
+    {"read_i32_le", (PyCFunction) read_i32_le, METH_FASTCALL,
+     PyDoc_STR("Read a 32-bit signed integer from bytes in little-endian format")
+    },
+    {"write_i64_le", (PyCFunction) write_i64_le, METH_FASTCALL,
+     PyDoc_STR("Write a 64-bit signed integer to BytesWriter in little-endian format")
+    },
+    {"read_i64_le", (PyCFunction) read_i64_le, METH_FASTCALL,
+     PyDoc_STR("Read a 64-bit signed integer from bytes in little-endian format")
+    },
 #endif
     {NULL, NULL, 0, NULL}
 };