[mypyc] Add vec extend function to librt.vecs (#21340)

This is similar to the `append` function, but the second parameter is an
iterable instead of a single item.

This is fairly straightforward, but we need multiple implementations for
the different vec representations. We also have separate functions for
`extend` with two vec arguments, and for `extend` with any iterable as
the second argument.

Also support the `extend(v, v)` special case.

We can later add optimized support for the buffer protocol (e.g. extend
`vec[u8]` with a `bytes` object).

Author: JukkaL

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

@@ -20,3 +20,4 @@ class vec(Generic[T]):
 def append(v: vec[T], o: T, /) -> vec[T]: ...
 def remove(v: vec[T], o: T, /) -> vec[T]: ...
 def pop(v: vec[T], i: i64 = -1, /) -> tuple[vec[T], T]: ...
+def extend(v: vec[T], o: Iterable[T], /) -> vec[T]: ...

mypyc/irbuild/specialize.py

@@ -99,7 +99,7 @@
     join_formatted_strings,
     tokenizer_format_call,
 )
-from mypyc.irbuild.vec import vec_append, vec_pop, vec_remove
+from mypyc.irbuild.vec import vec_append, vec_extend, vec_pop, vec_remove
 from mypyc.primitives.bytearray_ops import isinstance_bytearray
 from mypyc.primitives.bytes_ops import (
     bytes_adjust_index_op,
@@ -1520,6 +1520,19 @@ def translate_vec_append(builder: IRBuilder, expr: CallExpr, callee: RefExpr) ->
     return None
 
 
+@specialize_function("librt.vecs.extend")
+def translate_vec_extend(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
+    if len(expr.args) == 2 and expr.arg_kinds == [ARG_POS, ARG_POS]:
+        vec_arg = expr.args[0]
+        iter_arg = expr.args[1]
+        vec_type = builder.node_type(vec_arg)
+        if isinstance(vec_type, RVec):
+            vec_value = builder.accept(vec_arg)
+            iter_value = builder.accept(iter_arg)
+            return vec_extend(builder.builder, vec_value, iter_value, iter_arg.line)
+    return None
+
+
 @specialize_function("librt.vecs.remove")
 def translate_vec_remove(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
     if len(expr.args) == 2 and expr.arg_kinds == [ARG_POS, ARG_POS]:

mypyc/irbuild/vec.py

@@ -435,6 +435,38 @@ def vec_append(builder: LowLevelIRBuilder, vec: Value, item: Value, line: int) -
     return call
 
 
+def vec_extend(builder: LowLevelIRBuilder, vec: Value, iterable: Value, line: int) -> Value:
+    vec_type = vec.type
+    assert isinstance(vec_type, RVec)
+    item_type = vec_type.item_type
+    if isinstance(iterable.type, RVec) and iterable.type == vec_type:
+        suffix = "_vec"
+        src = iterable
+    else:
+        suffix = ""
+        src = builder.coerce(iterable, object_rprimitive, line)
+    item_type_arg: list[Value] = []
+    api_name = vec_api_by_item_type.get(item_type)
+    if api_name is not None:
+        name = f"{api_name}.extend{suffix}"
+    elif vec_type.depth() == 0:
+        name = f"VecTApi.extend{suffix}"
+        item_type_arg = [vec_item_type(builder, item_type, line)]
+    else:
+        name = f"VecNestedApi.extend{suffix}"
+    return builder.add(
+        CallC(
+            name,
+            [vec, src] + item_type_arg,
+            vec_type,
+            steals=[True, False] + ([False] if item_type_arg else []),
+            is_borrowed=False,
+            error_kind=ERR_MAGIC,
+            line=line,
+        )
+    )
+
+
 def vec_pop(builder: LowLevelIRBuilder, base: Value, index: Value, line: int) -> Value:
     assert isinstance(base.type, RVec)
     vec_type = base.type

mypyc/lib-rt/vecs/librt_vecs.c

@@ -876,6 +876,78 @@ static PyObject *vec_pop(PyObject *self, PyObject *args)
     return res;
 }
 
+static PyObject *vec_extend(PyObject *self, PyObject *args)
+{
+    PyObject *vec;
+    PyObject *iterable;
+
+    if (!PyArg_ParseTuple(args, "OO:extend", &vec, &iterable))
+        return NULL;
+
+    if (VecI64_Check(vec)) {
+        VecI64 v = ((VecI64Object *)vec)->vec;
+        VEC_INCREF(v);
+        v = VecI64_Extend(v, iterable);
+        if (VEC_IS_ERROR(v))
+            return NULL;
+        return VecI64_Box(v);
+    } else if (VecU8_Check(vec)) {
+        VecU8 v = ((VecU8Object *)vec)->vec;
+        VEC_INCREF(v);
+        v = VecU8_Extend(v, iterable);
+        if (VEC_IS_ERROR(v))
+            return NULL;
+        return VecU8_Box(v);
+    } else if (VecFloat_Check(vec)) {
+        VecFloat v = ((VecFloatObject *)vec)->vec;
+        VEC_INCREF(v);
+        v = VecFloat_Extend(v, iterable);
+        if (VEC_IS_ERROR(v))
+            return NULL;
+        return VecFloat_Box(v);
+    } else if (VecI32_Check(vec)) {
+        VecI32 v = ((VecI32Object *)vec)->vec;
+        VEC_INCREF(v);
+        v = VecI32_Extend(v, iterable);
+        if (VEC_IS_ERROR(v))
+            return NULL;
+        return VecI32_Box(v);
+    } else if (VecI16_Check(vec)) {
+        VecI16 v = ((VecI16Object *)vec)->vec;
+        VEC_INCREF(v);
+        v = VecI16_Extend(v, iterable);
+        if (VEC_IS_ERROR(v))
+            return NULL;
+        return VecI16_Box(v);
+    } else if (VecBool_Check(vec)) {
+        VecBool v = ((VecBoolObject *)vec)->vec;
+        VEC_INCREF(v);
+        v = VecBool_Extend(v, iterable);
+        if (VEC_IS_ERROR(v))
+            return NULL;
+        return VecBool_Box(v);
+    } else if (VecT_Check(vec)) {
+        VecT v = ((VecTObject *)vec)->vec;
+        size_t item_type = v.buf->item_type;
+        VEC_INCREF(v);
+        v = VecT_Extend(v, iterable, item_type);
+        if (VEC_IS_ERROR(v))
+            return NULL;
+        return VecT_Box(v, item_type);
+    } else if (VecNested_Check(vec)) {
+        VecNested v = ((VecNestedObject *)vec)->vec;
+        VEC_INCREF(v);
+        v = VecNested_Extend(v, iterable);
+        if (VEC_IS_ERROR(v))
+            return NULL;
+        return VecNested_Box(v);
+    } else {
+        PyErr_Format(PyExc_TypeError, "vec argument expected, got %.100s",
+                     Py_TYPE(vec)->tp_name);
+        return NULL;
+    }
+}
+
 // Return the base VecType for isinstance checks
 static PyTypeObject *get_vec_type(void) {
     return &VecType;
@@ -900,6 +972,7 @@ static PyMethodDef VecsMethods[] = {
     {"append",  vec_append, METH_VARARGS, "Append a value to the end of a vec"},
     {"remove",  vec_remove, METH_VARARGS, "Remove first occurrence of value from a vec"},
     {"pop",  vec_pop, METH_VARARGS, "Remove and return vec item at index (default last)"},
+    {"extend",  vec_extend, METH_VARARGS, "Extend a vec with items from an iterable"},
 #endif
     {NULL, NULL, 0, NULL}        /* Sentinel */
 };

mypyc/lib-rt/vecs/librt_vecs.h

@@ -278,10 +278,8 @@ typedef struct _VecI64API {
     VecI64 (*remove)(VecI64, int64_t);
     // TODO: Py_ssize_t
     VecI64 (*slice)(VecI64, int64_t, int64_t);
-    // PyObject *(*extend)(PyObject *, PyObject *);
-    // PyObject *(*concat)(PyObject *, PyObject *);
-    // bool (*contains)(PyObject *, int64_t);
-    // iter?
+    VecI64 (*extend)(VecI64, PyObject *);
+    VecI64 (*extend_vec)(VecI64, VecI64);
 } VecI64API;
 
 // vec[i32] operations + type objects (stored in a capsule)
@@ -297,10 +295,8 @@ typedef struct _VecI32API {
     VecI32 (*remove)(VecI32, int32_t);
     // TODO: Py_ssize_t
     VecI32 (*slice)(VecI32, int64_t, int64_t);
-    // PyObject *(*extend)(PyObject *, PyObject *);
-    // PyObject *(*concat)(PyObject *, PyObject *);
-    // bool (*contains)(PyObject *, int32_t);
-    // iter?
+    VecI32 (*extend)(VecI32, PyObject *);
+    VecI32 (*extend_vec)(VecI32, VecI32);
 } VecI32API;
 
 // vec[i16] operations + type objects (stored in a capsule)
@@ -316,10 +312,8 @@ typedef struct _VecI16API {
     VecI16 (*remove)(VecI16, int16_t);
     // TODO: Py_ssize_t
     VecI16 (*slice)(VecI16, int64_t, int64_t);
-    // PyObject *(*extend)(PyObject *, PyObject *);
-    // PyObject *(*concat)(PyObject *, PyObject *);
-    // bool (*contains)(PyObject *, int16_t);
-    // iter?
+    VecI16 (*extend)(VecI16, PyObject *);
+    VecI16 (*extend_vec)(VecI16, VecI16);
 } VecI16API;
 
 // vec[u8] operations + type objects (stored in a capsule)
@@ -335,10 +329,8 @@ typedef struct _VecU8API {
     VecU8 (*remove)(VecU8, uint8_t);
     // TODO: Py_ssize_t
     VecU8 (*slice)(VecU8, int64_t, int64_t);
-    // PyObject *(*extend)(PyObject *, PyObject *);
-    // PyObject *(*concat)(PyObject *, PyObject *);
-    // bool (*contains)(PyObject *, uint8_t);
-    // iter?
+    VecU8 (*extend)(VecU8, PyObject *);
+    VecU8 (*extend_vec)(VecU8, VecU8);
 } VecU8API;
 
 // vec[float] operations + type objects (stored in a capsule)
@@ -354,10 +346,8 @@ typedef struct _VecFloatAPI {
     VecFloat (*remove)(VecFloat, double);
     // TODO: Py_ssize_t
     VecFloat (*slice)(VecFloat, int64_t, int64_t);
-    // PyObject *(*extend)(PyObject *, PyObject *);
-    // PyObject *(*concat)(PyObject *, PyObject *);
-    // bool (*contains)(PyObject *, double);
-    // iter?
+    VecFloat (*extend)(VecFloat, PyObject *);
+    VecFloat (*extend_vec)(VecFloat, VecFloat);
 } VecFloatAPI;
 
 // vec[bool] operations + type objects (stored in a capsule)
@@ -373,10 +363,8 @@ typedef struct _VecBoolAPI {
     VecBool (*remove)(VecBool, char);
     // TODO: Py_ssize_t
     VecBool (*slice)(VecBool, int64_t, int64_t);
-    // PyObject *(*extend)(PyObject *, PyObject *);
-    // PyObject *(*concat)(PyObject *, PyObject *);
-    // bool (*contains)(PyObject *, char);
-    // iter?
+    VecBool (*extend)(VecBool, PyObject *);
+    VecBool (*extend_vec)(VecBool, VecBool);
 } VecBoolAPI;
 
 #ifndef MYPYC_DECLARED_tuple_T2VOO
@@ -405,10 +393,8 @@ typedef struct _VecTAPI {
     VecT (*remove)(VecT, PyObject *);
     // TODO: Py_ssize_t
     VecT (*slice)(VecT, int64_t, int64_t);
-    // PyObject *(*extend)(PyObject *, PyObject *);
-    // PyObject *(*concat)(PyObject *, PyObject *);
-    // bool (*contains)(PyObject *, PyObject *);
-    // iter?
+    VecT (*extend)(VecT, PyObject *, size_t);
+    VecT (*extend_vec)(VecT, VecT, size_t);
 } VecTAPI;
 
 
@@ -436,10 +422,8 @@ typedef struct _VecNestedAPI {
     VecNested (*remove)(VecNested, VecNestedBufItem);
     // TODO: Py_ssize_t
     VecNested (*slice)(VecNested, int64_t, int64_t);
-    // PyObject *(*extend)(PyObject *, PyObject *);
-    // PyObject *(*concat)(PyObject *, PyObject *);
-    // bool (*contains)(PyObject *, PyObject *);
-    // iter?
+    VecNested (*extend)(VecNested, PyObject *);
+    VecNested (*extend_vec)(VecNested, VecNested);
 } VecNestedAPI;
 
 typedef struct {
@@ -539,6 +523,8 @@ static inline int VecI64_IsUnboxError(int64_t x) {
 
 PyObject *VecI64_Box(VecI64);
 VecI64 VecI64_Append(VecI64, int64_t x);
+VecI64 VecI64_Extend(VecI64, PyObject *iterable);
+VecI64 VecI64_ExtendVec(VecI64 dst, VecI64 src);
 VecI64 VecI64_Remove(VecI64, int64_t x);
 VecI64PopResult VecI64_Pop(VecI64 v, Py_ssize_t index);
 
@@ -569,6 +555,8 @@ static inline int VecI32_IsUnboxError(int32_t x) {
 
 PyObject *VecI32_Box(VecI32);
 VecI32 VecI32_Append(VecI32, int32_t x);
+VecI32 VecI32_Extend(VecI32, PyObject *iterable);
+VecI32 VecI32_ExtendVec(VecI32 dst, VecI32 src);
 VecI32 VecI32_Remove(VecI32, int32_t x);
 VecI32PopResult VecI32_Pop(VecI32 v, Py_ssize_t index);
 
@@ -599,6 +587,8 @@ static inline int VecI16_IsUnboxError(int16_t x) {
 
 PyObject *VecI16_Box(VecI16);
 VecI16 VecI16_Append(VecI16, int16_t x);
+VecI16 VecI16_Extend(VecI16, PyObject *iterable);
+VecI16 VecI16_ExtendVec(VecI16 dst, VecI16 src);
 VecI16 VecI16_Remove(VecI16, int16_t x);
 VecI16PopResult VecI16_Pop(VecI16 v, Py_ssize_t index);
 
@@ -632,6 +622,8 @@ static inline int VecU8_IsUnboxError(uint8_t x) {
 
 PyObject *VecU8_Box(VecU8);
 VecU8 VecU8_Append(VecU8, uint8_t x);
+VecU8 VecU8_Extend(VecU8, PyObject *iterable);
+VecU8 VecU8_ExtendVec(VecU8 dst, VecU8 src);
 VecU8 VecU8_Remove(VecU8, uint8_t x);
 VecU8PopResult VecU8_Pop(VecU8 v, Py_ssize_t index);
 
@@ -655,6 +647,8 @@ static inline int VecFloat_IsUnboxError(double x) {
 
 PyObject *VecFloat_Box(VecFloat);
 VecFloat VecFloat_Append(VecFloat, double x);
+VecFloat VecFloat_Extend(VecFloat, PyObject *iterable);
+VecFloat VecFloat_ExtendVec(VecFloat dst, VecFloat src);
 VecFloat VecFloat_Remove(VecFloat, double x);
 VecFloatPopResult VecFloat_Pop(VecFloat v, Py_ssize_t index);
 
@@ -691,6 +685,8 @@ static inline int VecBool_IsUnboxError(char x) {
 
 PyObject *VecBool_Box(VecBool);
 VecBool VecBool_Append(VecBool, char x);
+VecBool VecBool_Extend(VecBool, PyObject *iterable);
+VecBool VecBool_ExtendVec(VecBool dst, VecBool src);
 VecBool VecBool_Remove(VecBool, char x);
 VecBoolPopResult VecBool_Pop(VecBool v, Py_ssize_t index);
 
@@ -716,6 +712,8 @@ VecT VecT_New(Py_ssize_t size, Py_ssize_t cap, size_t item_type);
 PyObject *VecT_FromIterable(size_t item_type, PyObject *iterable, int64_t cap);
 PyObject *VecT_Box(VecT vec, size_t item_type);
 VecT VecT_Append(VecT vec, PyObject *x, size_t item_type);
+VecT VecT_Extend(VecT vec, PyObject *iterable, size_t item_type);
+VecT VecT_ExtendVec(VecT dst, VecT src, size_t item_type);
 VecT VecT_Remove(VecT vec, PyObject *x);
 VecTPopResult VecT_Pop(VecT v, Py_ssize_t index);
 
@@ -729,6 +727,8 @@ VecNested VecNested_New(Py_ssize_t size, Py_ssize_t cap, size_t item_type, size_
 PyObject *VecNested_FromIterable(size_t item_type, size_t depth, PyObject *iterable, int64_t cap);
 PyObject *VecNested_Box(VecNested);
 VecNested VecNested_Append(VecNested vec, VecNestedBufItem x);
+VecNested VecNested_Extend(VecNested vec, PyObject *iterable);
+VecNested VecNested_ExtendVec(VecNested dst, VecNested src);
 VecNested VecNested_Remove(VecNested vec, VecNestedBufItem x);
 VecNestedPopResult VecNested_Pop(VecNested v, Py_ssize_t index);