[mypyc] Optimize list->vec and tuple->vec conversions (#21375)

Now we specialize these conversions for each vec variant (except for
nested, where it doesn't seem as useful). Previously they'd use append
for each item, which tends to be slow. They also used generic iterators
in some use cases, which is inefficient, especially for a small number
of items where the iterator construction and freeing is significant.

---------

Co-authored-by: Piotr Sawicki <sawickipiotr@outlook.com>

Author: JukkaL

mypyc/irbuild/expression.py

@@ -93,6 +93,7 @@
     is_list_rprimitive,
     is_none_rprimitive,
     is_object_rprimitive,
+    is_tuple_rprimitive,
     object_rprimitive,
     set_rprimitive,
     vec_api_by_item_type,
@@ -127,6 +128,7 @@
     vec_create,
     vec_create_from_values,
     vec_create_initialized,
+    vec_item_type,
     vec_slice,
 )
 from mypyc.primitives.bytes_ops import bytes_slice_op
@@ -630,26 +632,40 @@ def vec_from_iterable(
     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 (
+    use_c_from_iterable = (
         is_object_rprimitive(iterable_rtype)
+        or is_list_rprimitive(iterable_rtype)
+        or is_tuple_rprimitive(iterable_rtype)
+    )
+    if api_name is not None and (
+        use_c_from_iterable
         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.
+        name = f"{api_name}.from_iterable"
+        extra_args: list[Value] = []
+    elif api_name is None and vec_type.depth() == 0 and use_c_from_iterable:
+        name = "VecTApi.from_iterable"
+        extra_args = [vec_item_type(builder.builder, item_type, line)]
+    else:
+        name = None
+    if name is not None:
         iterable_val = builder.accept(iterable)
         cap = (
             as_platform_int(builder.builder, capacity, line)
             if capacity is not None
             else Integer(0, int64_rprimitive)
         )
+        args = extra_args + [iterable_val, cap]
         call = CallC(
-            f"{api_name}.from_iterable",
-            [iterable_val, cap],
+            name,
+            args,
             vec_type,
-            steals=[False, False],
+            steals=[False] * len(args),
             is_borrowed=False,
             error_kind=ERR_MAGIC,
             line=line,

mypyc/lib-rt/vecs/librt_vecs.c

@@ -114,7 +114,10 @@ static PyObject *vec_generic_alias_call(PyObject *self, PyObject *args, PyObject
                 return NULL;
             return VecT_Box(vec, p->item_type);
         } else {
-            return VecT_FromIterable(p->item_type, init, cap);
+            VecT vec = VecT_FromIterable(p->item_type, init, cap);
+            if (VEC_IS_ERROR(vec))
+                return NULL;
+            return VecT_Box(vec, p->item_type);
         }
     } else {
         if (init == NULL) {

mypyc/lib-rt/vecs/librt_vecs.h

@@ -400,6 +400,7 @@ typedef struct _VecTAPI {
     VecT (*remove)(VecT, PyObject *);
     // TODO: Py_ssize_t
     VecT (*slice)(VecT, int64_t, int64_t);
+    VecT (*from_iterable)(size_t, PyObject *, int64_t);
     VecT (*extend)(VecT, PyObject *, size_t);
     VecT (*extend_vec)(VecT, VecT, size_t);
 } VecTAPI;
@@ -716,7 +717,7 @@ static inline int VecT_ItemCheck(VecT v, PyObject *item, size_t item_type) {
 }
 
 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);
+VecT 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);

mypyc/lib-rt/vecs/vec_t.c

@@ -685,10 +685,41 @@ PyTypeObject VecTType = {
     // TODO: free
 };
 
-PyObject *VecT_FromIterable(size_t item_type, PyObject *iterable, int64_t cap) {
+static inline VecT vec_from_sequence(
+        size_t item_type, PyObject *seq, int64_t cap, const int is_list) {
+    Py_ssize_t n = is_list ? PyList_GET_SIZE(seq) : PyTuple_GET_SIZE(seq);
+    Py_ssize_t alloc_size = n > cap ? n : cap;
+    VecT v = vec_alloc(alloc_size, item_type);
+    if (VEC_IS_ERROR(v))
+        return vec_error();
+    for (Py_ssize_t i = 0; i < n; i++) {
+        PyObject *item = is_list ? PyList_GET_ITEM(seq, i) : PyTuple_GET_ITEM(seq, i);
+        if (!VecT_ItemCheck(v, item, item_type)) {
+            for (Py_ssize_t j = i; j < alloc_size; j++)
+                v.buf->items[j] = NULL;
+            VEC_DECREF(v);
+            return vec_error();
+        }
+        Py_INCREF(item);
+        v.buf->items[i] = item;
+    }
+    for (Py_ssize_t j = n; j < alloc_size; j++)
+        v.buf->items[j] = NULL;
+    vec_track_buffer(&v);
+    v.len = n;
+    return v;
+}
+
+VecT VecT_FromIterable(size_t item_type, PyObject *iterable, int64_t cap) {
+    if (PyList_CheckExact(iterable)) {
+        return vec_from_sequence(item_type, iterable, cap, 1);
+    } else if (PyTuple_CheckExact(iterable)) {
+        return vec_from_sequence(item_type, iterable, cap, 0);
+    }
+
     VecT v = vec_alloc(cap, item_type);
     if (VEC_IS_ERROR(v))
-        return NULL;
+        return vec_error();
     if (cap > 0) {
         for (int64_t i = 0; i < cap; i++)
             v.buf->items[i] = NULL;
@@ -699,30 +730,30 @@ PyObject *VecT_FromIterable(size_t item_type, 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) {
         if (!VecT_ItemCheck(v, item, item_type)) {
             Py_DECREF(iter);
             VEC_DECREF(v);
             Py_DECREF(item);
-            return NULL;
+            return vec_error();
         }
         v = VecT_Append(v, item, item_type);
         Py_DECREF(item);
         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 VecT_Box(v, item_type);
+    return v;
 }
 
 VecTAPI Vec_TAPI = {
@@ -736,6 +767,7 @@ VecTAPI Vec_TAPI = {
     VecT_Pop,
     VecT_Remove,
     VecT_Slice,
+    VecT_FromIterable,
     VecT_Extend,
     VecT_ExtendVec,
 };

mypyc/lib-rt/vecs/vec_template.c

@@ -135,6 +135,25 @@ inline static int vec_get_buffer(PyObject *obj, Py_buffer *view) {
 }
 #endif
 
+static inline VEC vec_from_sequence(PyObject *seq, int64_t cap, const int is_list) {
+    Py_ssize_t n = is_list ? PyList_GET_SIZE(seq) : PyTuple_GET_SIZE(seq);
+    Py_ssize_t alloc_size = n > cap ? n : cap;
+    VEC v = vec_alloc(alloc_size);
+    if (VEC_IS_ERROR(v))
+        return vec_error();
+    for (Py_ssize_t i = 0; i < n; i++) {
+        PyObject *item = is_list ? PyList_GET_ITEM(seq, i) : PyTuple_GET_ITEM(seq, i);
+        ITEM_C_TYPE x = UNBOX_ITEM(item);
+        if (IS_UNBOX_ERROR(x)) {
+            VEC_DECREF(v);
+            return vec_error();
+        }
+        v.buf->items[i] = x;
+    }
+    v.len = n;
+    return v;
+}
+
 VEC FUNC(FromIterable)(PyObject *iterable, int64_t cap) {
     if (cap < 0) {
         PyErr_SetString(PyExc_ValueError, "capacity must not be negative");
@@ -175,6 +194,12 @@ VEC FUNC(FromIterable)(PyObject *iterable, int64_t cap) {
     }
 #endif
 
+    if (PyList_CheckExact(iterable)) {
+        return vec_from_sequence(iterable, cap, 1);
+    } else if (PyTuple_CheckExact(iterable)) {
+        return vec_from_sequence(iterable, cap, 0);
+    }
+
     VEC v = vec_alloc(cap);
     if (VEC_IS_ERROR(v))
         return vec_error();

mypyc/test-data/irbuild-vec-i64.test

@@ -943,3 +943,31 @@ def f(x):
 L0:
     r0 = VecI64Api.from_iterable(x, 0)
     return r0
+
+[case testVecI64CreateFromListObject]
+from librt.vecs import vec
+from mypy_extensions import i64
+
+def from_list(a: list[i64]) -> vec[i64]:
+    return vec[i64](a)
+[out]
+def from_list(a):
+    a :: list
+    r0 :: vec[i64]
+L0:
+    r0 = VecI64Api.from_iterable(a, 0)
+    return r0
+
+[case testVecI64CreateFromTupleObject]
+from librt.vecs import vec
+from mypy_extensions import i64
+
+def from_tuple(a: tuple[i64, ...]) -> vec[i64]:
+    return vec[i64](a)
+[out]
+def from_tuple(a):
+    a :: tuple
+    r0 :: vec[i64]
+L0:
+    r0 = VecI64Api.from_iterable(a, 0)
+    return r0

mypyc/test-data/irbuild-vec-t.test

@@ -453,3 +453,54 @@ L0:
     r8 = r7 + 8
     keep_alive r4
     return r4
+
+[case testVecTConstructFromIterable]
+from typing import Iterable
+from librt.vecs import vec
+
+def f(x: Iterable[str]) -> vec[str]:
+    return vec[str](x)
+[out]
+def f(x):
+    x, r0 :: object
+    r1 :: ptr
+    r2 :: vec[str]
+L0:
+    r0 = load_address PyUnicode_Type
+    r1 = r0
+    r2 = VecTApi.from_iterable(r1, x, 0)
+    return r2
+
+[case testVecTCreateFromListObject]
+from librt.vecs import vec
+
+def from_list(a: list[str]) -> vec[str]:
+    return vec[str](a)
+[out]
+def from_list(a):
+    a :: list
+    r0 :: object
+    r1 :: ptr
+    r2 :: vec[str]
+L0:
+    r0 = load_address PyUnicode_Type
+    r1 = r0
+    r2 = VecTApi.from_iterable(r1, a, 0)
+    return r2
+
+[case testVecTCreateFromTupleObject]
+from librt.vecs import vec
+
+def from_tuple(a: tuple[str, ...]) -> vec[str]:
+    return vec[str](a)
+[out]
+def from_tuple(a):
+    a :: tuple
+    r0 :: object
+    r1 :: ptr
+    r2 :: vec[str]
+L0:
+    r0 = load_address PyUnicode_Type
+    r1 = r0
+    r2 = VecTApi.from_iterable(r1, a, 0)
+    return r2

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

@@ -44,6 +44,40 @@ def test_construct_from_initializer() -> None:
     with assertRaises(TypeError):
         vec[i64]([None])
 
+def test_construct_from_tuple() -> None:
+    v = vec[i64](())
+    assert len(v) == 0
+    v = vec[i64]((1, 2, 3))
+    assert len(v) == 3
+    assert v[0] == 1
+    assert v[1] == 2
+    assert v[2] == 3
+    v = vec[i64]((2**63 - 1, -2**63))
+    assert len(v) == 2
+    assert v[0] == 2**63 - 1
+    assert v[1] == -2**63
+    with assertRaises(TypeError):
+        vec[i64](('x',))
+    with assertRaises(TypeError):
+        vec[i64]((None,))
+    with assertRaises(TypeError):
+        vec[i64]((1, 'x', 2))
+
+def test_construct_from_tuple_with_capacity() -> None:
+    v = vec[i64]((1, 2), capacity=5)
+    assert len(v) == 2
+    assert v[0] == 1
+    assert v[1] == 2
+    old = v
+    v = append(v, 3)
+    v = append(v, 4)
+    v = append(v, 5)
+    old[0] = 99
+    assert v[0] == 99
+    v = vec[i64]((1, 2, 3), capacity=1)
+    assert len(v) == 3
+    assert v[2] == 3
+
 def test_append() -> None:
     v = vec[i64]()
     v = append(v, 1)

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

@@ -99,6 +99,13 @@ def test_basic_int_operations() -> None:
         with assertRaises(TypeError):
             append(v, "0")
 
+def test_construct_bool_from_list() -> None:
+    v = vec[bool]([True, False, True])
+    assert len(v) == 3
+    assert v[0] is True
+    assert v[1] is False
+    assert v[2] is True
+
 def test_equality() -> None:
     for t1 in ITEM_TYPES:
         for t2 in ITEM_TYPES:

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

@@ -63,6 +63,17 @@ def test_construct_from_initializer_nested() -> None:
     with assertRaises(TypeError):
         vec[str]([vec[bytes]()])
 
+def test_construct_from_tuple_nested() -> None:
+    v = vec[vec[str]](())
+    assert len(v) == 0
+    v = vec[vec[str]]((vec[str](['a']), vec[str](['b', 'c'])))
+    assert len(v) == 2
+    assert v[0][0] == 'a'
+    assert v[1][0] == 'b'
+    assert v[1][1] == 'c'
+    with assertRaises(TypeError):
+        vec[vec[str]]((vec[str](), 'x'))
+
 def test_isinstance() -> None:
     assert isinstance(vec[vec[str]](), vec)
     assert isinstance(vec[vec[Optional[str]]](), vec)