[mypyc] Use memory copy for vec init and extend using buffer protocol (#21349)

For vecs with value item types, use the buffer protocol for fast bulk
item access during construction and `extend()`. Previously each item had
to be boxed and then unboxed, which was inefficient.

This works with `bytes` and `array.array`, for example, and should also
work with NumPy arrays (didn't test with NumPy though). Examples where
this matters include `v = vec[u8](b"foo")` and `v = extend(v, b"bar")`.

Since the `bool` item type uses a non-standard encoding, it's not
supported.

We can make `bytes` even faster by adding further specialization. I'll
do this in a follow-up PR.

I used coding agent assist with small incremental changes.

Author: JukkaL

mypyc/irbuild/expression.py

@@ -62,8 +62,10 @@
 from mypyc.ir.class_ir import ClassIR
 from mypyc.ir.func_ir import FUNC_CLASSMETHOD, FUNC_STATICMETHOD
 from mypyc.ir.ops import (
+    ERR_MAGIC,
     Assign,
     BasicBlock,
+    CallC,
     ComparisonOp,
     Integer,
     LoadAddress,
@@ -80,8 +82,11 @@
     RTuple,
     RVec,
     bool_rprimitive,
+    int64_rprimitive,
     int_rprimitive,
     is_any_int,
+    is_bytearray_rprimitive,
+    is_bytes_rprimitive,
     is_fixed_width_rtype,
     is_int64_rprimitive,
     is_int_rprimitive,
@@ -90,6 +95,7 @@
     is_object_rprimitive,
     object_rprimitive,
     set_rprimitive,
+    vec_api_by_item_type,
 )
 from mypyc.irbuild.ast_helpers import is_borrow_friendly_expr, process_conditional
 from mypyc.irbuild.builder import IRBuilder, int_borrow_friendly_op
@@ -116,6 +122,7 @@
     translate_object_setattr,
 )
 from mypyc.irbuild.vec import (
+    as_platform_int,
     vec_append,
     vec_create,
     vec_create_from_values,
@@ -620,9 +627,37 @@ def vec_from_iterable(
     capacity: Value | None = None,
 ) -> Value:
     """Construct a vec from an arbitrary iterable."""
-    # Translate it as a vec comprehension vec[t]([<name> for <name> in
-    # iterable]). This way we can use various special casing supported
-    # by for loops and comprehensions.
+    item_type = vec_type.item_type
+    api_name = vec_api_by_item_type.get(item_type)
+    iterable_rtype = builder.node_type(iterable)
+    if api_name is not None and (
+        is_object_rprimitive(iterable_rtype)
+        or is_bytes_rprimitive(iterable_rtype)
+        or is_bytearray_rprimitive(iterable_rtype)
+    ):
+        # For generic iterables (typed as object) and bytes/bytearray
+        # (which support the buffer protocol for fast memcpy), call the
+        # C-level from_iterable. For concrete types like range, list,
+        # vec, etc., the for-loop desugaring below produces better IR.
+        iterable_val = builder.accept(iterable)
+        cap = (
+            as_platform_int(builder.builder, capacity, line)
+            if capacity is not None
+            else Integer(0, int64_rprimitive)
+        )
+        call = CallC(
+            f"{api_name}.from_iterable",
+            [iterable_val, cap],
+            vec_type,
+            steals=[False, False],
+            is_borrowed=False,
+            error_kind=ERR_MAGIC,
+            line=line,
+        )
+        return builder.add(call)
+
+    # Use a for loop with vec_append. The comprehension helper
+    # special-cases range, list, vec, etc. for efficient iteration.
     vec = Register(vec_type)
     builder.assign(vec, vec_create(builder.builder, vec_type, 0, line, capacity=capacity), line)
     name = f"___tmp_{line}"

mypyc/lib-rt/vecs/librt_vecs.h

@@ -278,6 +278,7 @@ typedef struct _VecI64API {
     VecI64 (*remove)(VecI64, int64_t);
     // TODO: Py_ssize_t
     VecI64 (*slice)(VecI64, int64_t, int64_t);
+    VecI64 (*from_iterable)(PyObject *, int64_t);
     VecI64 (*extend)(VecI64, PyObject *);
     VecI64 (*extend_vec)(VecI64, VecI64);
 } VecI64API;
@@ -295,6 +296,7 @@ typedef struct _VecI32API {
     VecI32 (*remove)(VecI32, int32_t);
     // TODO: Py_ssize_t
     VecI32 (*slice)(VecI32, int64_t, int64_t);
+    VecI32 (*from_iterable)(PyObject *, int64_t);
     VecI32 (*extend)(VecI32, PyObject *);
     VecI32 (*extend_vec)(VecI32, VecI32);
 } VecI32API;
@@ -312,6 +314,7 @@ typedef struct _VecI16API {
     VecI16 (*remove)(VecI16, int16_t);
     // TODO: Py_ssize_t
     VecI16 (*slice)(VecI16, int64_t, int64_t);
+    VecI16 (*from_iterable)(PyObject *, int64_t);
     VecI16 (*extend)(VecI16, PyObject *);
     VecI16 (*extend_vec)(VecI16, VecI16);
 } VecI16API;
@@ -329,6 +332,7 @@ typedef struct _VecU8API {
     VecU8 (*remove)(VecU8, uint8_t);
     // TODO: Py_ssize_t
     VecU8 (*slice)(VecU8, int64_t, int64_t);
+    VecU8 (*from_iterable)(PyObject *, int64_t);
     VecU8 (*extend)(VecU8, PyObject *);
     VecU8 (*extend_vec)(VecU8, VecU8);
 } VecU8API;
@@ -346,6 +350,7 @@ typedef struct _VecFloatAPI {
     VecFloat (*remove)(VecFloat, double);
     // TODO: Py_ssize_t
     VecFloat (*slice)(VecFloat, int64_t, int64_t);
+    VecFloat (*from_iterable)(PyObject *, int64_t);
     VecFloat (*extend)(VecFloat, PyObject *);
     VecFloat (*extend_vec)(VecFloat, VecFloat);
 } VecFloatAPI;
@@ -363,6 +368,7 @@ typedef struct _VecBoolAPI {
     VecBool (*remove)(VecBool, char);
     // TODO: Py_ssize_t
     VecBool (*slice)(VecBool, int64_t, int64_t);
+    VecBool (*from_iterable)(PyObject *, int64_t);
     VecBool (*extend)(VecBool, PyObject *);
     VecBool (*extend_vec)(VecBool, VecBool);
 } VecBoolAPI;

mypyc/lib-rt/vecs/vec_float.c

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

mypyc/lib-rt/vecs/vec_i16.c

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

mypyc/lib-rt/vecs/vec_i32.c

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

mypyc/lib-rt/vecs/vec_i64.c

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

mypyc/lib-rt/vecs/vec_template.c

@@ -100,10 +100,72 @@ VEC FUNC(New)(Py_ssize_t size, Py_ssize_t cap) {
     return vec;
 }
 
-PyObject *FUNC(FromIterable)(PyObject *iterable, int64_t cap) {
+#ifdef BUFFER_FORMAT_CHAR_OK
+inline static int buffer_format_matches(const char *fmt) {
+    char c = *fmt;
+    if (c == '@' || c == '=') {
+        c = fmt[1];
+    }
+#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
+    else if (c == '<') { c = fmt[1]; }
+    else if (c == '>' || c == '!') { return 0; }
+#else
+    else if (c == '>') { c = fmt[1]; }
+    else if (c == '<' || c == '!') { return 0; }
+#endif
+    return c != '\0' && BUFFER_FORMAT_CHAR_OK(c);
+}
+
+// Try to get a compatible buffer view from 'obj'. Return 1 if successful
+// (view is filled and caller must call PyBuffer_Release), 0 if the object
+// doesn't support buffer protocol or the format doesn't match (no cleanup
+// needed), or -1 on error.
+inline static int vec_get_buffer(PyObject *obj, Py_buffer *view) {
+    if (PyObject_GetBuffer(obj, view, PyBUF_C_CONTIGUOUS | PyBUF_FORMAT) != 0) {
+        PyErr_Clear();
+        return 0;
+    }
+    if (view->ndim == 1
+        && view->itemsize == sizeof(ITEM_C_TYPE)
+        && buffer_format_matches(view->format)) {
+        return 1;
+    }
+    PyBuffer_Release(view);
+    return 0;
+}
+#endif
+
+VEC FUNC(FromIterable)(PyObject *iterable, int64_t cap) {
+    if (cap < 0) {
+        PyErr_SetString(PyExc_ValueError, "capacity must not be negative");
+        return vec_error();
+    }
+
+#ifdef BUFFER_FORMAT_CHAR_OK
+    Py_buffer view;
+    int buf_ok = vec_get_buffer(iterable, &view);
+    if (buf_ok < 0)
+        return vec_error();
+    if (buf_ok) {
+        Py_ssize_t n = view.len / (Py_ssize_t)sizeof(ITEM_C_TYPE);
+        Py_ssize_t alloc_size = n > cap ? n : cap;
+        VEC v = vec_alloc(alloc_size);
+        if (VEC_IS_ERROR(v)) {
+            PyBuffer_Release(&view);
+            return vec_error();
+        }
+        if (n > 0) {
+            memcpy(v.buf->items, view.buf, n * sizeof(ITEM_C_TYPE));
+        }
+        v.len = n;
+        PyBuffer_Release(&view);
+        return v;
+    }
+#endif
+
     VEC v = vec_alloc(cap);
     if (VEC_IS_ERROR(v))
-        return NULL;
+        return vec_error();
     if (cap > 0) {
         memset(v.buf->items, 0, sizeof(ITEM_C_TYPE) * cap);
     }
@@ -112,7 +174,7 @@ PyObject *FUNC(FromIterable)(PyObject *iterable, int64_t cap) {
     PyObject *iter = PyObject_GetIter(iterable);
     if (iter == NULL) {
         VEC_DECREF(v);
-        return NULL;
+        return vec_error();
     }
     PyObject *item;
     while ((item = PyIter_Next(iter)) != NULL) {
@@ -121,21 +183,21 @@ PyObject *FUNC(FromIterable)(PyObject *iterable, int64_t cap) {
         if (IS_UNBOX_ERROR(x)) {
             Py_DECREF(iter);
             VEC_DECREF(v);
-            return NULL;
+            return vec_error();
         }
         v = FUNC(Append)(v, x);
         if (VEC_IS_ERROR(v)) {
             Py_DECREF(iter);
             VEC_DECREF(v);
-            return NULL;
+            return vec_error();
         }
     }
     Py_DECREF(iter);
     if (PyErr_Occurred()) {
         VEC_DECREF(v);
-        return NULL;
+        return vec_error();
     }
-    return FUNC(Box)(v);
+    return v;
 }
 
 static PyObject *vec_new(PyTypeObject *self, PyObject *args, PyObject *kw) {
@@ -152,7 +214,10 @@ static PyObject *vec_new(PyTypeObject *self, PyObject *args, PyObject *kw) {
     if (init == NULL) {
         return FUNC(Box)(FUNC(New)(0, cap));
     } else {
-        return (PyObject *)FUNC(FromIterable)(init, cap);
+        VEC v = FUNC(FromIterable)(init, cap);
+        if (VEC_IS_ERROR(v))
+            return NULL;
+        return FUNC(Box)(v);
     }
 }
 
@@ -332,13 +397,68 @@ VEC FUNC(Append)(VEC vec, ITEM_C_TYPE x) {
     }
 }
 
+// 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.
+inline static VEC vec_extend_items(
+    VEC dst, const ITEM_C_TYPE *items, Py_ssize_t n, int force_alloc
+) {
+    if (unlikely(n > PY_SSIZE_T_MAX - dst.len)) {
+        PyErr_NoMemory();
+        VEC_DECREF(dst);
+        return vec_error();
+    }
+    Py_ssize_t new_len = dst.len + n;
+    Py_ssize_t cap = dst.buf ? VEC_CAP(dst) : 0;
+    if (!force_alloc && new_len <= cap) {
+        memcpy(dst.buf->items + dst.len, items, sizeof(ITEM_C_TYPE) * n);
+        dst.len = new_len;
+        return dst;
+    }
+    Py_ssize_t new_cap = cap;
+    while (new_cap < new_len) {
+        if (unlikely(new_cap > (PY_SSIZE_T_MAX - 1) / 2)) {
+            new_cap = new_len;
+            break;
+        }
+        new_cap = 2 * new_cap + 1;
+    }
+    VEC new = vec_alloc(new_cap);
+    if (VEC_IS_ERROR(new)) {
+        VEC_DECREF(dst);
+        return vec_error();
+    }
+    if (dst.len > 0)
+        memcpy(new.buf->items, dst.buf->items, sizeof(ITEM_C_TYPE) * dst.len);
+    memcpy(new.buf->items + dst.len, items, sizeof(ITEM_C_TYPE) * n);
+    new.len = new_len;
+    Py_XDECREF(dst.buf);
+    return new;
+}
+
 // Extend 'vec' with items from 'iterable', stealing 'vec'.
 // Return extended 'vec', or error vec on failure.
 VEC FUNC(Extend)(VEC vec, PyObject *iterable) {
     if (Py_TYPE(iterable) == &VEC_TYPE) {
         return FUNC(ExtendVec)(vec, ((VEC_OBJECT *)iterable)->vec);
     }
 
+#ifdef BUFFER_FORMAT_CHAR_OK
+    Py_buffer view;
+    int buf_ok = vec_get_buffer(iterable, &view);
+    if (buf_ok < 0) {
+        VEC_DECREF(vec);
+        return vec_error();
+    }
+    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);
+        PyBuffer_Release(&view);
+        return vec;
+    }
+#endif
+
     PyObject *iter = PyObject_GetIter(iterable);
     if (iter == NULL) {
         VEC_DECREF(vec);
@@ -372,39 +492,7 @@ VEC FUNC(Extend)(VEC vec, PyObject *iterable) {
 VEC FUNC(ExtendVec)(VEC dst, VEC src) {
     if (src.len == 0)
         return dst;
-    if (unlikely(src.len > PY_SSIZE_T_MAX - dst.len)) {
-        PyErr_NoMemory();
-        VEC_DECREF(dst);
-        return vec_error();
-    }
-    Py_ssize_t new_len = dst.len + src.len;
-    Py_ssize_t cap = dst.buf ? VEC_CAP(dst) : 0;
-    if (new_len <= cap && dst.buf != src.buf) {
-        // Fast path: enough capacity and no aliasing
-        memcpy(dst.buf->items + dst.len, src.buf->items, sizeof(ITEM_C_TYPE) * src.len);
-        dst.len = new_len;
-        return dst;
-    }
-    // Need to reallocate (or dst and src share a buffer)
-    Py_ssize_t new_cap = cap;
-    while (new_cap < new_len) {
-        if (unlikely(new_cap > (PY_SSIZE_T_MAX - 1) / 2)) {
-            new_cap = new_len;
-            break;
-        }
-        new_cap = 2 * new_cap + 1;
-    }
-    VEC new = vec_alloc(new_cap);
-    if (VEC_IS_ERROR(new)) {
-        VEC_DECREF(dst);
-        return vec_error();
-    }
-    if (dst.len > 0)
-        memcpy(new.buf->items, dst.buf->items, sizeof(ITEM_C_TYPE) * dst.len);
-    memcpy(new.buf->items + dst.len, src.buf->items, sizeof(ITEM_C_TYPE) * src.len);
-    new.len = new_len;
-    Py_XDECREF(dst.buf);
-    return new;
+    return vec_extend_items(dst, src.buf->items, src.len, dst.buf == src.buf);
 }
 
 // Remove item from 'vec', stealing 'vec'. Return 'vec' with item removed.
@@ -575,6 +663,7 @@ NAME(API) FEATURES = {
     FUNC(Pop),
     FUNC(Remove),
     FUNC(Slice),
+    FUNC(FromIterable),
     FUNC(Extend),
     FUNC(ExtendVec),
 };