Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
JukkaL merged 14 commits into from
Apr 27, 2026
Merged

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

Show file tree
Hide file tree
Changes from 13 commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
cc7975e
Add librt.vecs.extend
JukkaL Apr 8, 2026
0c20457
Add extend primitive
JukkaL Apr 8, 2026
134d1f0
WIP extend
JukkaL Apr 13, 2026
fb49c84
Add comment
JukkaL Apr 23, 2026
33283b9
Add tests
JukkaL Apr 23, 2026
43b3bcf
Fix issues with extending with itself
JukkaL Apr 23, 2026
91351f3
Fix aliasing issue
JukkaL Apr 24, 2026
1223fc0
Dispatch based on second param type in wrapper function
JukkaL Apr 27, 2026
45e0054
Fix test case
JukkaL Apr 27, 2026
b83f1a5
Lint
JukkaL Apr 27, 2026
cf02530
Update comment
JukkaL Apr 27, 2026
8d75f7c
Perform specialization in API function instead of wrapper
JukkaL Apr 27, 2026
c5510a5
Fix integer overflows
JukkaL Apr 27, 2026
85dad0a
Address feedback
JukkaL Apr 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions mypy/typeshed/stubs/librt/librt/vecs.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]: ...
15 changes: 14 additions & 1 deletion mypyc/irbuild/specialize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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]:
Expand Down
58 changes: 58 additions & 0 deletions mypyc/irbuild/vec.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -435,6 +435,64 @@ 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
# Fast path: extend vec with another vec of the same type (no boxing needed)
if isinstance(iterable.type, RVec) and iterable.type == vec_type:
return _vec_extend_vec(builder, vec, iterable, line)
Copy link
Copy Markdown
Collaborator

@p-sawicki p-sawicki Apr 27, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the fast path is almost exactly the same so we could fold it into this function to reduce duplication

if isinstance(iterable.type, RVec) and iterable.type == vec_type:
   suffix = "_vec"
   src = iterable
else:
   suffix = ""
   src = builder.coerce(iterable, object_rprimitive, line)

and then call extend{suffix}

# Slow path: extend with generic iterable
coerced_iterable = 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"
elif vec_type.depth() == 0:
name = "VecTApi.extend"
item_type_arg = [vec_item_type(builder, item_type, line)]
else:
name = "VecNestedApi.extend"
return builder.add(
CallC(
name,
[vec, coerced_iterable] + 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_extend_vec(builder: LowLevelIRBuilder, vec: Value, src: Value, line: int) -> Value:
"""Extend vec with another vec of the same type (fast path, no boxing)."""
vec_type = vec.type
assert isinstance(vec_type, RVec)
item_type = vec_type.item_type
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_vec"
elif vec_type.depth() == 0:
name = "VecTApi.extend_vec"
item_type_arg = [vec_item_type(builder, item_type, line)]
else:
name = "VecNestedApi.extend_vec"
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
Expand Down
73 changes: 73 additions & 0 deletions mypyc/lib-rt/vecs/librt_vecs.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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 */
};
Expand Down
64 changes: 32 additions & 32 deletions mypyc/lib-rt/vecs/librt_vecs.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -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;


Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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);

Expand All @@ -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);

Expand Down Expand Up @@ -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);

Expand All @@ -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);

Expand All @@ -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);

Expand Down
Loading
Loading