[mypyc] Add buffer protocol support for vec (#21359)

Now `vec` with value item types can be used as 1-dimensional buffers.
This allows
other code access raw memory in the vec buffer.

Currently only read-only buffer access is supported.

I relied on coding agent assist quite a bit.

Author: JukkaL

mypy/typeshed/stubs/librt/librt/vecs.pyi

@@ -16,6 +16,7 @@ class vec(Generic[T]):
     def __setitem__(self, i: i64, o: T, /) -> None: ...
     def __contains__(self, o: object, /) -> bool: ...
     def __iter__(self) -> Iterator[T]: ...
+    def __buffer__(self, flags: int, /) -> memoryview: ...
 
 def append(v: vec[T], o: T, /) -> vec[T]: ...
 def remove(v: vec[T], o: T, /) -> vec[T]: ...

mypyc/lib-rt/vecs/vec_bool.c

@@ -14,6 +14,7 @@
 #define BOX_ITEM VecBool_BoxItem
 #define UNBOX_ITEM VecBool_UnboxItem
 #define IS_UNBOX_ERROR VecBool_IsUnboxError
+#define BUFFER_FORMAT "b"
 
 #include "vec_template.c"
 

mypyc/lib-rt/vecs/vec_float.c

@@ -15,6 +15,7 @@
 #define UNBOX_ITEM VecFloat_UnboxItem
 #define IS_UNBOX_ERROR VecFloat_IsUnboxError
 #define BUFFER_FORMAT_CHAR_OK(c) ((c) == 'd')
+#define BUFFER_FORMAT "d"
 
 #include "vec_template.c"
 

mypyc/lib-rt/vecs/vec_i16.c

@@ -15,6 +15,7 @@
 #define UNBOX_ITEM VecI16_UnboxItem
 #define IS_UNBOX_ERROR VecI16_IsUnboxError
 #define BUFFER_FORMAT_CHAR_OK(c) ((c) == 'h')
+#define BUFFER_FORMAT "h"
 
 #include "vec_template.c"
 

mypyc/lib-rt/vecs/vec_i32.c

@@ -15,6 +15,7 @@
 #define UNBOX_ITEM VecI32_UnboxItem
 #define IS_UNBOX_ERROR VecI32_IsUnboxError
 #define BUFFER_FORMAT_CHAR_OK(c) ((c) == 'i' || ((c) == 'l' && sizeof(long) == 4))
+#define BUFFER_FORMAT "i"
 
 #include "vec_template.c"
 

mypyc/lib-rt/vecs/vec_i64.c

@@ -15,6 +15,7 @@
 #define UNBOX_ITEM VecI64_UnboxItem
 #define IS_UNBOX_ERROR VecI64_IsUnboxError
 #define BUFFER_FORMAT_CHAR_OK(c) ((c) == 'q' || ((c) == 'l' && sizeof(long) == 8))
+#define BUFFER_FORMAT "q"
 
 #include "vec_template.c"
 

mypyc/lib-rt/vecs/vec_template.c

@@ -409,6 +409,16 @@ VEC FUNC(Append)(VEC vec, ITEM_C_TYPE x) {
     }
 }
 
+inline static int vec_memory_overlaps(const void *p1, Py_ssize_t len1,
+                                      const void *p2, Py_ssize_t len2) {
+    if (len1 <= 0 || len2 <= 0)
+        return 0;
+    uintptr_t a = (uintptr_t)p1, b = (uintptr_t)p2;
+    if (a <= b)
+        return b - a < (uintptr_t)len1;
+    return a - b < (uintptr_t)len2;
+}
+
 // Extend 'dst' by appending 'n' items from 'items', stealing 'dst'.
 // Caller guarantees n > 0 and that 'items' remains valid for the call.
 // If force_alloc is true, always allocate a new buffer even when dst has capacity.
@@ -471,8 +481,14 @@ VEC FUNC(Extend)(VEC vec, PyObject *iterable) {
     }
     if (buf_ok) {
         Py_ssize_t n = view.len / (Py_ssize_t)sizeof(ITEM_C_TYPE);
-        if (n > 0)
-            vec = vec_extend_items(vec, (const ITEM_C_TYPE *)view.buf, n, 0);
+        if (n > 0) {
+            Py_ssize_t dst_bytes = n * (Py_ssize_t)sizeof(ITEM_C_TYPE);
+            int force_alloc = vec.buf != NULL
+                && n <= VEC_CAP(vec) - vec.len
+                && vec_memory_overlaps(view.buf, view.len,
+                                       vec.buf->items + vec.len, dst_bytes);
+            vec = vec_extend_items(vec, (const ITEM_C_TYPE *)view.buf, n, force_alloc);
+        }
         PyBuffer_Release(&view);
         return vec;
     }
@@ -566,6 +582,46 @@ static PyMappingMethods vec_mapping_methods = {
     .mp_subscript = vec_subscript,
 };
 
+#ifdef BUFFER_FORMAT
+static int vec_getbuffer(VEC_OBJECT *self, Py_buffer *view, int flags) {
+    if (view == NULL) {
+        PyErr_SetString(PyExc_BufferError,
+            "vec_getbuffer: view==NULL argument is obsolete");
+        return -1;
+    }
+    if ((flags & PyBUF_WRITABLE) == PyBUF_WRITABLE) {
+        PyErr_SetString(PyExc_BufferError, "Object is not writable");
+        view->obj = NULL;
+        return -1;
+    }
+
+    view->obj = (PyObject *)self;
+    Py_INCREF(self);
+    view->buf = (self->vec.buf != NULL) ? (void *)self->vec.buf->items : NULL;
+    view->len = self->vec.len * (Py_ssize_t)sizeof(ITEM_C_TYPE);
+    view->readonly = 1;
+    view->itemsize = sizeof(ITEM_C_TYPE);
+    view->format = NULL;
+    if ((flags & PyBUF_FORMAT) == PyBUF_FORMAT)
+        view->format = BUFFER_FORMAT;
+    view->ndim = 1;
+    view->shape = NULL;
+    if ((flags & PyBUF_ND) == PyBUF_ND)
+        view->shape = &self->vec.len;
+    view->strides = NULL;
+    if ((flags & PyBUF_STRIDES) == PyBUF_STRIDES)
+        view->strides = &view->itemsize;
+    view->suboffsets = NULL;
+    view->internal = NULL;
+
+    return 0;
+}
+
+static PyBufferProcs vec_buffer_procs = {
+    .bf_getbuffer = (getbufferproc)vec_getbuffer,
+};
+#endif
+
 static PySequenceMethods vec_sequence_methods = {
     .sq_item = vec_get_item,
     .sq_ass_item = vec_ass_item,
@@ -667,6 +723,9 @@ PyTypeObject VEC_TYPE = {
     .tp_iter = vec_iter,
     .tp_as_sequence = &vec_sequence_methods,
     .tp_as_mapping = &vec_mapping_methods,
+#ifdef BUFFER_FORMAT
+    .tp_as_buffer = &vec_buffer_procs,
+#endif
     .tp_richcompare = vec_richcompare,
     .tp_methods = vec_methods,
 };

mypyc/lib-rt/vecs/vec_u8.c

@@ -15,6 +15,7 @@
 #define UNBOX_ITEM VecU8_UnboxItem
 #define IS_UNBOX_ERROR VecU8_IsUnboxError
 #define BUFFER_FORMAT_CHAR_OK(c) ((c) == 'B' || (c) == 'c')
+#define BUFFER_FORMAT "B"
 
 #include "vec_template.c"
 

mypyc/test-data/run-vecs-misc-interp.test

@@ -531,6 +531,108 @@ def test_extend_bool() -> None:
     v = extend(v, [False, True])
     assert v == vec[bool]([True, False, True])
 
+[case testLibrtVecsBufferProtocol_librt_experimental]
+import struct
+from typing import Any
+
+import librt.vecs
+import mypy_extensions
+
+from testutil import assertRaises
+
+vec: Any = librt.vecs.vec
+append: Any = librt.vecs.append
+extend: Any = librt.vecs.extend
+i64: Any = mypy_extensions.i64
+i32: Any = mypy_extensions.i32
+i16: Any = mypy_extensions.i16
+u8: Any = mypy_extensions.u8
+memoryview_: Any = eval('memoryview')
+
+def test_buffer_u8() -> None:
+    v = vec[u8]([1, 2, 3])
+    mv = memoryview_(v)
+    assert mv.format == "B"
+    assert mv.itemsize == 1
+    assert mv.readonly
+    assert mv.ndim == 1
+    assert mv.shape == (3,)
+    assert list(mv) == [1, 2, 3]
+
+def test_buffer_i64() -> None:
+    v = vec[i64]([10, 20, -30])
+    mv = memoryview_(v)
+    assert mv.format == "q"
+    assert mv.itemsize == 8
+    assert mv.readonly
+    assert mv.ndim == 1
+    assert mv.shape == (3,)
+    assert list(mv) == [10, 20, -30]
+
+def test_buffer_i32() -> None:
+    v = vec[i32]([5, -12])
+    mv = memoryview_(v)
+    assert mv.format == "i"
+    assert mv.itemsize == 4
+    assert mv.readonly
+    assert list(mv) == [5, -12]
+
+def test_buffer_i16() -> None:
+    v = vec[i16]([100, -200])
+    mv = memoryview_(v)
+    assert mv.format == "h"
+    assert mv.itemsize == 2
+    assert mv.readonly
+    assert list(mv) == [100, -200]
+
+def test_buffer_float() -> None:
+    v = vec[float]([1.5, -2.5])
+    mv = memoryview_(v)
+    assert mv.format == "d"
+    assert mv.itemsize == 8
+    assert mv.readonly
+    assert list(mv) == [1.5, -2.5]
+
+def test_buffer_bool() -> None:
+    v = vec[bool]([True, False, True])
+    mv = memoryview_(v)
+    assert mv.format == "b"
+    assert mv.itemsize == 1
+    assert mv.readonly
+    assert list(mv) == [1, 0, 1]
+
+def test_buffer_empty() -> None:
+    v = vec[i64]([])
+    mv = memoryview_(v)
+    assert mv.format == "q"
+    assert len(mv) == 0
+    assert mv.shape == (0,)
+    assert list(mv) == []
+
+def test_buffer_readonly() -> None:
+    v = vec[u8]([1, 2, 3])
+    mv = memoryview_(v)
+    with assertRaises(TypeError):
+        mv[0] = 99
+
+def test_buffer_bytes_conversion() -> None:
+    v = vec[u8]([72, 101, 108, 108, 111])
+    assert bytes(v) == b"Hello"
+
+def test_buffer_tobytes() -> None:
+    v = vec[i32]([1, 2])
+    mv = memoryview_(v)
+    b = mv.tobytes()
+    values = struct.unpack("ii", b)
+    assert values == (1, 2)
+
+def test_buffer_extend_from_shared_buffer() -> None:
+    old = vec[u8]([1, 2], capacity=5)
+    new = append(old, 3)
+    v = extend(old, memoryview_(new))
+    assert v == vec[u8]([1, 2, 1, 2, 3])
+    assert new == vec[u8]([1, 2, 3])
+
 [case testLibrtVecsFeaturesNotAvailableInNonExperimentalBuild_librt]
 # This also ensures librt.vecs can be built without experimental features
 import librt.vecs