From b30b4894e15c4a60e8a91c97497250161f59d718 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Thu, 5 Feb 2026 14:39:10 +0000 Subject: [PATCH 1/5] WIP iterator support --- mypyc/lib-rt/vecs/librt_vecs.c | 16 ++++ mypyc/lib-rt/vecs/librt_vecs.h | 10 +++ mypyc/lib-rt/vecs/vec_nested.c | 82 +++++++++++++++++++ mypyc/lib-rt/vecs/vec_t.c | 83 +++++++++++++++++++ mypyc/lib-rt/vecs/vec_template.c | 60 ++++++++++++++ mypyc/test-data/run-vecs-i64-interp.test | 39 +++++++++ mypyc/test-data/run-vecs-misc-interp.test | 80 ++++++++++++++++++ mypyc/test-data/run-vecs-nested-interp.test | 89 +++++++++++++++++++++ mypyc/test-data/run-vecs-t-interp.test | 46 +++++++++++ 9 files changed, 505 insertions(+) diff --git a/mypyc/lib-rt/vecs/librt_vecs.c b/mypyc/lib-rt/vecs/librt_vecs.c index 3696e696d5c3e..13d540a1dc451 100644 --- a/mypyc/lib-rt/vecs/librt_vecs.c +++ b/mypyc/lib-rt/vecs/librt_vecs.c @@ -950,36 +950,52 @@ librt_vecs_module_exec(PyObject *m) return -1; if (PyType_Ready(&VecTBufType) < 0) return -1; + if (PyType_Ready(&VecTIterType) < 0) + return -1; if (PyType_Ready(&VecNestedType) < 0) return -1; if (PyType_Ready(&VecNestedBufType) < 0) return -1; + if (PyType_Ready(&VecNestedIterType) < 0) + return -1; if (PyType_Ready(&VecI64Type) < 0) return -1; if (PyType_Ready(&VecI64BufType) < 0) return -1; + if (PyType_Ready(&VecI64IterType) < 0) + return -1; if (PyType_Ready(&VecI32Type) < 0) return -1; if (PyType_Ready(&VecI32BufType) < 0) return -1; + if (PyType_Ready(&VecI32IterType) < 0) + return -1; if (PyType_Ready(&VecI16Type) < 0) return -1; if (PyType_Ready(&VecI16BufType) < 0) return -1; + if (PyType_Ready(&VecI16IterType) < 0) + return -1; if (PyType_Ready(&VecU8Type) < 0) return -1; if (PyType_Ready(&VecU8BufType) < 0) return -1; + if (PyType_Ready(&VecU8IterType) < 0) + return -1; if (PyType_Ready(&VecFloatType) < 0) return -1; if (PyType_Ready(&VecFloatBufType) < 0) return -1; + if (PyType_Ready(&VecFloatIterType) < 0) + return -1; if (PyType_Ready(&VecBoolType) < 0) return -1; if (PyType_Ready(&VecBoolBufType) < 0) return -1; + if (PyType_Ready(&VecBoolIterType) < 0) + return -1; Py_INCREF(&VecType); if (PyModule_AddObject(m, "vec", (PyObject *)&VecType) < 0) { diff --git a/mypyc/lib-rt/vecs/librt_vecs.h b/mypyc/lib-rt/vecs/librt_vecs.h index 80dc977c19146..e9b93157c7982 100644 --- a/mypyc/lib-rt/vecs/librt_vecs.h +++ b/mypyc/lib-rt/vecs/librt_vecs.h @@ -493,6 +493,16 @@ extern PyTypeObject VecBoolType; extern PyTypeObject VecTType; extern PyTypeObject VecNestedType; +// Iterator type objects for vec iteration +extern PyTypeObject VecI64IterType; +extern PyTypeObject VecI32IterType; +extern PyTypeObject VecI16IterType; +extern PyTypeObject VecU8IterType; +extern PyTypeObject VecFloatIterType; +extern PyTypeObject VecBoolIterType; +extern PyTypeObject VecTIterType; +extern PyTypeObject VecNestedIterType; + // Type objects corresponding to the 'i64', 'i32', 'i16, and 'u8' types extern PyTypeObject *LibRTVecs_I64TypeObj; extern PyTypeObject *LibRTVecs_I32TypeObj; diff --git a/mypyc/lib-rt/vecs/vec_nested.c b/mypyc/lib-rt/vecs/vec_nested.c index 0034d9e9bfff5..75b777803011a 100644 --- a/mypyc/lib-rt/vecs/vec_nested.c +++ b/mypyc/lib-rt/vecs/vec_nested.c @@ -420,6 +420,87 @@ static PyMethodDef vec_methods[] = { {NULL, NULL, 0, NULL}, /* Sentinel */ }; +// Iterator type for nested vecs + +typedef struct { + PyObject_HEAD + VecNestedObject *vec_obj; // Reference to vec object (keeps buffer alive) + Py_ssize_t index; // Current iteration index +} VecNestedIterObject; + +PyTypeObject VecNestedIterType; + +static PyObject *VecNested_iter(PyObject *self) { + VecNestedIterObject *it = PyObject_GC_New(VecNestedIterObject, &VecNestedIterType); + if (it == NULL) + return NULL; + Py_INCREF(self); + it->vec_obj = (VecNestedObject *)self; + it->index = 0; + PyObject_GC_Track(it); + return (PyObject *)it; +} + +static int +VecNestedIter_traverse(VecNestedIterObject *self, visitproc visit, void *arg) +{ + Py_VISIT(self->vec_obj); + return 0; +} + +static int +VecNestedIter_clear(VecNestedIterObject *self) +{ + Py_CLEAR(self->vec_obj); + return 0; +} + +static void VecNestedIter_dealloc(VecNestedIterObject *self) { + PyObject_GC_UnTrack(self); + Py_XDECREF(self->vec_obj); + PyObject_GC_Del(self); +} + +static PyObject *VecNestedIter_next(VecNestedIterObject *self) { + if (self->vec_obj == NULL) + return NULL; + VecNested v = self->vec_obj->vec; + if (self->index < v.len) { + PyObject *item = box_vec_item_by_index(v, self->index); + self->index++; + return item; + } + return NULL; // StopIteration +} + +static PyObject *VecNestedIter_len(VecNestedIterObject *self, PyObject *Py_UNUSED(ignored)) { + if (self->vec_obj == NULL) + return PyLong_FromSsize_t(0); + Py_ssize_t remaining = self->vec_obj->vec.len - self->index; + if (remaining < 0) + remaining = 0; + return PyLong_FromSsize_t(remaining); +} + +static PyMethodDef VecNestedIter_methods[] = { + {"__length_hint__", (PyCFunction)VecNestedIter_len, METH_NOARGS, NULL}, + {NULL, NULL, 0, NULL}, +}; + +PyTypeObject VecNestedIterType = { + PyVarObject_HEAD_INIT(NULL, 0) + .tp_name = "vec_nested_iterator", + .tp_basicsize = sizeof(VecNestedIterObject), + .tp_itemsize = 0, + .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, + .tp_traverse = (traverseproc)VecNestedIter_traverse, + .tp_clear = (inquiry)VecNestedIter_clear, + .tp_dealloc = (destructor)VecNestedIter_dealloc, + .tp_iter = PyObject_SelfIter, + .tp_iternext = (iternextfunc)VecNestedIter_next, + .tp_methods = VecNestedIter_methods, +}; + PyTypeObject VecNestedBufType = { PyVarObject_HEAD_INIT(NULL, 0) .tp_name = "vecbuf", @@ -447,6 +528,7 @@ PyTypeObject VecNestedType = { .tp_dealloc = (destructor)VecNested_dealloc, //.tp_free = PyObject_GC_Del, .tp_repr = (reprfunc)vec_repr, + .tp_iter = VecNested_iter, .tp_as_sequence = &VecNestedSequence, .tp_as_mapping = &VecNestedMapping, .tp_richcompare = vec_richcompare, diff --git a/mypyc/lib-rt/vecs/vec_t.c b/mypyc/lib-rt/vecs/vec_t.c index 8a03ea1a2e644..c85eec42be70e 100644 --- a/mypyc/lib-rt/vecs/vec_t.c +++ b/mypyc/lib-rt/vecs/vec_t.c @@ -416,6 +416,88 @@ static PyMethodDef vec_methods[] = { {NULL, NULL, 0, NULL}, /* Sentinel */ }; +// Iterator type for vec[T] (reference types) + +typedef struct { + PyObject_HEAD + VecTObject *vec_obj; // Reference to vec object (keeps buffer alive) + Py_ssize_t index; // Current iteration index +} VecTIterObject; + +PyTypeObject VecTIterType; + +static PyObject *VecT_iter(PyObject *self) { + VecTIterObject *it = PyObject_GC_New(VecTIterObject, &VecTIterType); + if (it == NULL) + return NULL; + Py_INCREF(self); + it->vec_obj = (VecTObject *)self; + it->index = 0; + PyObject_GC_Track(it); + return (PyObject *)it; +} + +static int +VecTIter_traverse(VecTIterObject *self, visitproc visit, void *arg) +{ + Py_VISIT(self->vec_obj); + return 0; +} + +static int +VecTIter_clear(VecTIterObject *self) +{ + Py_CLEAR(self->vec_obj); + return 0; +} + +static void VecTIter_dealloc(VecTIterObject *self) { + PyObject_GC_UnTrack(self); + Py_XDECREF(self->vec_obj); + PyObject_GC_Del(self); +} + +static PyObject *VecTIter_next(VecTIterObject *self) { + if (self->vec_obj == NULL) + return NULL; + VecT v = self->vec_obj->vec; + if (self->index < v.len) { + PyObject *item = v.buf->items[self->index]; + self->index++; + Py_INCREF(item); + return item; + } + return NULL; // StopIteration +} + +static PyObject *VecTIter_len(VecTIterObject *self, PyObject *Py_UNUSED(ignored)) { + if (self->vec_obj == NULL) + return PyLong_FromSsize_t(0); + Py_ssize_t remaining = self->vec_obj->vec.len - self->index; + if (remaining < 0) + remaining = 0; + return PyLong_FromSsize_t(remaining); +} + +static PyMethodDef VecTIter_methods[] = { + {"__length_hint__", (PyCFunction)VecTIter_len, METH_NOARGS, NULL}, + {NULL, NULL, 0, NULL}, +}; + +PyTypeObject VecTIterType = { + PyVarObject_HEAD_INIT(NULL, 0) + .tp_name = "vec_iterator", + .tp_basicsize = sizeof(VecTIterObject), + .tp_itemsize = 0, + .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, + .tp_traverse = (traverseproc)VecTIter_traverse, + .tp_clear = (inquiry)VecTIter_clear, + .tp_dealloc = (destructor)VecTIter_dealloc, + .tp_iter = PyObject_SelfIter, + .tp_iternext = (iternextfunc)VecTIter_next, + .tp_methods = VecTIter_methods, +}; + PyTypeObject VecTBufType = { PyVarObject_HEAD_INIT(NULL, 0) .tp_name = "vecbuf", @@ -443,6 +525,7 @@ PyTypeObject VecTType = { .tp_dealloc = (destructor)VecT_dealloc, //.tp_free = PyObject_GC_Del, .tp_repr = (reprfunc)vec_repr, + .tp_iter = VecT_iter, .tp_as_sequence = &VecTSequence, .tp_as_mapping = &VecTMapping, .tp_richcompare = vec_richcompare, diff --git a/mypyc/lib-rt/vecs/vec_template.c b/mypyc/lib-rt/vecs/vec_template.c index dc1e9df3609fc..cd74381688e51 100644 --- a/mypyc/lib-rt/vecs/vec_template.c +++ b/mypyc/lib-rt/vecs/vec_template.c @@ -354,6 +354,65 @@ static PyMethodDef vec_methods[] = { {NULL, NULL, 0, NULL}, /* Sentinel */ }; +// Iterator type for specialized vec types + +typedef struct { + PyObject_HEAD + VEC_OBJECT *vec_obj; // Reference to vec object (keeps buffer alive) + Py_ssize_t index; // Current iteration index +} NAME(IterObject); + +PyTypeObject NAME(IterType); + +static PyObject *vec_iter(PyObject *self) { + NAME(IterObject) *it = PyObject_New(NAME(IterObject), &NAME(IterType)); + if (it == NULL) + return NULL; + Py_INCREF(self); + it->vec_obj = (VEC_OBJECT *)self; + it->index = 0; + return (PyObject *)it; +} + +static void vec_iter_dealloc(NAME(IterObject) *self) { + Py_DECREF(self->vec_obj); + PyObject_Del(self); +} + +static PyObject *vec_iter_next(NAME(IterObject) *self) { + VEC v = self->vec_obj->vec; + if (self->index < v.len) { + PyObject *item = BOX_ITEM(v.buf->items[self->index]); + self->index++; + return item; + } + return NULL; // StopIteration +} + +static PyObject *vec_iter_len(NAME(IterObject) *self, PyObject *Py_UNUSED(ignored)) { + Py_ssize_t remaining = self->vec_obj->vec.len - self->index; + if (remaining < 0) + remaining = 0; + return PyLong_FromSsize_t(remaining); +} + +static PyMethodDef vec_iter_methods[] = { + {"__length_hint__", (PyCFunction)vec_iter_len, METH_NOARGS, NULL}, + {NULL, NULL, 0, NULL}, +}; + +PyTypeObject NAME(IterType) = { + PyVarObject_HEAD_INIT(NULL, 0) + .tp_name = "vec_iterator[" ITEM_TYPE_STR "]", + .tp_basicsize = sizeof(NAME(IterObject)), + .tp_itemsize = 0, + .tp_flags = Py_TPFLAGS_DEFAULT, + .tp_dealloc = (destructor)vec_iter_dealloc, + .tp_iter = PyObject_SelfIter, + .tp_iternext = (iternextfunc)vec_iter_next, + .tp_methods = vec_iter_methods, +}; + PyTypeObject BUF_TYPE = { PyVarObject_HEAD_INIT(NULL, 0) .tp_name = "vecbuf[" ITEM_TYPE_STR "]", @@ -377,6 +436,7 @@ PyTypeObject VEC_TYPE = { //.tp_free = PyObject_Del, .tp_dealloc = (destructor)vec_dealloc, .tp_repr = (reprfunc)vec_repr, + .tp_iter = vec_iter, .tp_as_sequence = &vec_sequence_methods, .tp_as_mapping = &vec_mapping_methods, .tp_richcompare = vec_richcompare, diff --git a/mypyc/test-data/run-vecs-i64-interp.test b/mypyc/test-data/run-vecs-i64-interp.test index 52d816881b66b..166d6818a724c 100644 --- a/mypyc/test-data/run-vecs-i64-interp.test +++ b/mypyc/test-data/run-vecs-i64-interp.test @@ -231,6 +231,45 @@ def test_iteration_nested() -> None: (5, 2), (5, 5), (5, 7), (7, 2), (7, 5), (7, 7)] +def test_iterator_protocol() -> None: + # Test iter() and next() functions explicitly + v = vec[i64]([1, 2, 3]) + it = iter(v) + assert next(it) == 1 + assert next(it) == 2 + assert next(it) == 3 + with assertRaises(StopIteration): + next(it) + + # Test multiple iterators on same vec + v = vec[i64]([10, 20]) + it1 = iter(v) + it2 = iter(v) + assert next(it1) == 10 + assert next(it2) == 10 + assert next(it1) == 20 + assert next(it2) == 20 + + # Test empty vec iterator + it = iter(vec[i64]()) + with assertRaises(StopIteration): + next(it) + +def test_iterator_length_hint() -> None: + v = vec[i64]([1, 2, 3, 4, 5]) + it = iter(v) + op: Any = getattr(sys, "modules")["operator"] + length_hint = op.length_hint + assert length_hint(it) == 5 + next(it) + assert length_hint(it) == 4 + next(it) + next(it) + assert length_hint(it) == 2 + next(it) + next(it) + assert length_hint(it) == 0 + def test_slicing() -> None: v = vec[i64](range(5)) assert v[1:4] == vec[i64]([1, 2, 3]) diff --git a/mypyc/test-data/run-vecs-misc-interp.test b/mypyc/test-data/run-vecs-misc-interp.test index db2f79d926608..07ba78c3f4a29 100644 --- a/mypyc/test-data/run-vecs-misc-interp.test +++ b/mypyc/test-data/run-vecs-misc-interp.test @@ -373,6 +373,86 @@ def test_bool_pop() -> None: assert n is True assert v == vec[bool]([True, False]) +# Iteration tests for all item types + +def test_iteration() -> None: + # Test iteration for each primitive type + for x in vec[float](): + assert False + a = [] + for x in vec[float]([1.5, 2.5, 3.5]): + a.append(x) + assert a == [1.5, 2.5, 3.5] + + for x in vec[u8](): + assert False + a = [] + for x in vec[u8]([1, 2, 255]): + a.append(x) + assert a == [1, 2, 255] + + for x in vec[i16](): + assert False + a = [] + for x in vec[i16]([-1000, 0, 1000]): + a.append(x) + assert a == [-1000, 0, 1000] + + for x in vec[i32](): + assert False + a = [] + for x in vec[i32]([-100000, 0, 100000]): + a.append(x) + assert a == [-100000, 0, 100000] + + for x in vec[bool](): + assert False + a = [] + for x in vec[bool]([True, False, True]): + a.append(x) + assert a == [True, False, True] + +def test_iterator_protocol() -> None: + # Test iter() and next() for float + v_float = vec[float]([1.5, 2.5]) + it = iter(v_float) + assert next(it) == 1.5 + assert next(it) == 2.5 + with assertRaises(StopIteration): + next(it) + + # Test iter() and next() for u8 + v_u8 = vec[u8]([10, 20]) + it = iter(v_u8) + assert next(it) == 10 + assert next(it) == 20 + with assertRaises(StopIteration): + next(it) + + # Test iter() and next() for i16 + v_i16 = vec[i16]([-100, 100]) + it = iter(v_i16) + assert next(it) == -100 + assert next(it) == 100 + with assertRaises(StopIteration): + next(it) + + # Test iter() and next() for i32 + v_i32 = vec[i32]([-1000, 1000]) + it = iter(v_i32) + assert next(it) == -1000 + assert next(it) == 1000 + with assertRaises(StopIteration): + next(it) + + # Test bool specifically since values are True/False singletons + v_bool = vec[bool]([True, False]) + it = iter(v_bool) + assert next(it) is True + assert next(it) is False + with assertRaises(StopIteration): + next(it) + [case testLibrtVecsFeaturesNotAvailableInNonExperimentalBuild_librt] # This also ensures librt.vecs can be built without experimental features import librt.vecs diff --git a/mypyc/test-data/run-vecs-nested-interp.test b/mypyc/test-data/run-vecs-nested-interp.test index a8ed7b87c426c..bd621e99498d1 100644 --- a/mypyc/test-data/run-vecs-nested-interp.test +++ b/mypyc/test-data/run-vecs-nested-interp.test @@ -321,3 +321,92 @@ def test_pop_index() -> None: v, item = pop(v, 0) assert item == vec[str](['15']) assert v == vec[vec[str]]() + +def test_iteration() -> None: + # Test empty nested vec iteration + for x in vec[vec[str]](): + assert False + + # Test single item iteration + a = [] + for x in vec[vec[str]]([vec[str](['a'])]): + a.append(x) + assert a == [vec[str](['a'])] + + # Test multiple items iteration + a = [] + for x in vec[vec[str]]([vec[str](['a']), vec[str](['b', 'c']), vec[str]()]): + a.append(x) + assert a == [vec[str](['a']), vec[str](['b', 'c']), vec[str]()] + + # Test iteration with vec[vec[i64]] + a = [] + for x in vec[vec[i64]]([vec[i64]([1, 2]), vec[i64]([3])]): + a.append(x) + assert a == [vec[i64]([1, 2]), vec[i64]([3])] + +def test_iteration_nested_loop() -> None: + # Test nested iteration over same vec + v = vec[vec[str]]([vec[str](['a']), vec[str](['b'])]) + a = [] + for x in v: + for y in v: + a.append((x, y)) + assert a == [ + (vec[str](['a']), vec[str](['a'])), + (vec[str](['a']), vec[str](['b'])), + (vec[str](['b']), vec[str](['a'])), + (vec[str](['b']), vec[str](['b'])), + ] + +def test_iteration_deeply_nested() -> None: + # Test vec[vec[vec[str]]] iteration + vvv = vec[vec[vec[str]]]([ + vec[vec[str]]([vec[str](['a'])]), + vec[vec[str]]([vec[str](['b']), vec[str](['c'])]) + ]) + a = [] + for x in vvv: + a.append(x) + assert len(a) == 2 + assert a[0] == vec[vec[str]]([vec[str](['a'])]) + assert a[1] == vec[vec[str]]([vec[str](['b']), vec[str](['c'])]) + +def test_iterator_protocol() -> None: + # Test iter() and next() functions explicitly + v = vec[vec[str]]([vec[str](['a']), vec[str](['b']), vec[str](['c'])]) + it = iter(v) + assert next(it) == vec[str](['a']) + assert next(it) == vec[str](['b']) + assert next(it) == vec[str](['c']) + with assertRaises(StopIteration): + next(it) + + # Test multiple iterators on same vec + v = vec[vec[i64]]([vec[i64]([1]), vec[i64]([2])]) + it1 = iter(v) + it2 = iter(v) + assert next(it1) == vec[i64]([1]) + assert next(it2) == vec[i64]([1]) + assert next(it1) == vec[i64]([2]) + assert next(it2) == vec[i64]([2]) + + # Test empty vec iterator + it = iter(vec[vec[str]]()) + with assertRaises(StopIteration): + next(it) + +def test_iterator_length_hint() -> None: + v = vec[vec[str]]([vec[str](['a']), vec[str](['b']), vec[str](['c']), vec[str](['d']), vec[str](['e'])]) + it = iter(v) + op: Any = getattr(sys, "modules")["operator"] + length_hint = op.length_hint + assert length_hint(it) == 5 + next(it) + assert length_hint(it) == 4 + next(it) + next(it) + assert length_hint(it) == 2 + next(it) + next(it) + assert length_hint(it) == 0 diff --git a/mypyc/test-data/run-vecs-t-interp.test b/mypyc/test-data/run-vecs-t-interp.test index cb527b50f7c24..881530b50f7a6 100644 --- a/mypyc/test-data/run-vecs-t-interp.test +++ b/mypyc/test-data/run-vecs-t-interp.test @@ -324,6 +324,52 @@ def test_iteration() -> None: a.append(x) assert a == ['11', None, '45', '64'] +def test_iterator_protocol() -> None: + # Test iter() and next() functions explicitly + v = vec[str](['a', 'b', 'c']) + it = iter(v) + assert next(it) == 'a' + assert next(it) == 'b' + assert next(it) == 'c' + with assertRaises(StopIteration): + next(it) + + # Test multiple iterators on same vec + v = vec[str](['x', 'y']) + it1 = iter(v) + it2 = iter(v) + assert next(it1) == 'x' + assert next(it2) == 'x' + assert next(it1) == 'y' + assert next(it2) == 'y' + + # Test empty vec iterator + it = iter(vec[str]()) + with assertRaises(StopIteration): + next(it) + + # Test iteration with optional type + v = vec[Optional[str]](['a', None, 'b']) + it = iter(v) + assert next(it) == 'a' + assert next(it) is None + assert next(it) == 'b' + +def test_iterator_length_hint() -> None: + v = vec[str](['a', 'b', 'c', 'd', 'e']) + it = iter(v) + op: Any = getattr(sys, "modules")["operator"] + length_hint = op.length_hint + assert length_hint(it) == 5 + next(it) + assert length_hint(it) == 4 + next(it) + next(it) + assert length_hint(it) == 2 + next(it) + next(it) + assert length_hint(it) == 0 + def test_slicing() -> None: v = vec[str](str(i) for i in range(5)) assert v[1:4] == vec[str](['1', '2', '3']) From c2f2dafa7a0609ed840b47caf2fa4c902554bc43 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Wed, 1 Apr 2026 10:21:11 +0100 Subject: [PATCH 2/5] More tests --- mypyc/test-data/run-vecs-i64.test | 23 +++++++++++++++++++++++ mypyc/test-data/run-vecs-nested.test | 12 ++++++++++++ mypyc/test-data/run-vecs-t.test | 9 +++++++++ 3 files changed, 44 insertions(+) diff --git a/mypyc/test-data/run-vecs-i64.test b/mypyc/test-data/run-vecs-i64.test index 1c19bfb11e767..0c1fb2f83e548 100644 --- a/mypyc/test-data/run-vecs-i64.test +++ b/mypyc/test-data/run-vecs-i64.test @@ -427,6 +427,29 @@ def test_pop_index() -> None: assert n == 15 assert v == vec[i64]() +def test_iter_and_next() -> None: + v = vec[i64]([1, 2, 3]) + it = iter(v) + assert next(it) == 1 + assert next(it) == 2 + assert next(it) == 3 + with assertRaises(StopIteration): + next(it) + + # Multiple iterators on same vec + v = vec[i64]([10, 20]) + it1 = iter(v) + it2 = iter(v) + assert next(it1) == 10 + assert next(it2) == 10 + assert next(it1) == 20 + assert next(it2) == 20 + + # Empty vec iterator + it = iter(vec[i64]()) + with assertRaises(StopIteration): + next(it) + def f(x: vec[i64]) -> None: pass diff --git a/mypyc/test-data/run-vecs-nested.test b/mypyc/test-data/run-vecs-nested.test index a8c8032121e7f..0cf430d9c06cc 100644 --- a/mypyc/test-data/run-vecs-nested.test +++ b/mypyc/test-data/run-vecs-nested.test @@ -418,6 +418,18 @@ def test_doubly_nested_vec_i64_operations() -> None: # TODO: box/unbox +def test_iter_and_next() -> None: + a = vec[str](["s1" + str()]) + b = vec[str](["s2" + str()]) + c = vec[str](["s3" + str()]) + v = vec[vec[str]]([a, b, c]) + it = iter(v) + assert next(it) == a + assert next(it) == b + assert next(it) == c + with assertRaises(StopIteration): + next(it) + def f(x: vec[vec[str]]) -> None: pass diff --git a/mypyc/test-data/run-vecs-t.test b/mypyc/test-data/run-vecs-t.test index 153428148518e..2b69e31408201 100644 --- a/mypyc/test-data/run-vecs-t.test +++ b/mypyc/test-data/run-vecs-t.test @@ -352,6 +352,15 @@ def test_pop_optional() -> None: v, s = pop(v, -1) assert s is None +def test_iter_and_next() -> None: + v = vec[str](["s1" + str(), "s2" + str(), "s3" + str()]) + it = iter(v) + assert next(it) == "s1" + assert next(it) == "s2" + assert next(it) == "s3" + with assertRaises(StopIteration): + next(it) + def f(x: vec[str]) -> None: pass From 682f37db9ea805070a462136fbb6fbbdec8750f1 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Wed, 1 Apr 2026 11:03:33 +0100 Subject: [PATCH 3/5] Fix error handling --- mypyc/lib-rt/vecs/vec_nested.c | 2 ++ mypyc/lib-rt/vecs/vec_template.c | 2 ++ 2 files changed, 4 insertions(+) diff --git a/mypyc/lib-rt/vecs/vec_nested.c b/mypyc/lib-rt/vecs/vec_nested.c index 75b777803011a..568190810012e 100644 --- a/mypyc/lib-rt/vecs/vec_nested.c +++ b/mypyc/lib-rt/vecs/vec_nested.c @@ -467,6 +467,8 @@ static PyObject *VecNestedIter_next(VecNestedIterObject *self) { VecNested v = self->vec_obj->vec; if (self->index < v.len) { PyObject *item = box_vec_item_by_index(v, self->index); + if (item == NULL) + return NULL; self->index++; return item; } diff --git a/mypyc/lib-rt/vecs/vec_template.c b/mypyc/lib-rt/vecs/vec_template.c index cd74381688e51..c21d4af9a4934 100644 --- a/mypyc/lib-rt/vecs/vec_template.c +++ b/mypyc/lib-rt/vecs/vec_template.c @@ -383,6 +383,8 @@ static PyObject *vec_iter_next(NAME(IterObject) *self) { VEC v = self->vec_obj->vec; if (self->index < v.len) { PyObject *item = BOX_ITEM(v.buf->items[self->index]); + if (item == NULL) + return NULL; self->index++; return item; } From 21cef985cbd9d16654e5fe03d3bb18f9d505fe35 Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Wed, 1 Apr 2026 11:11:59 +0100 Subject: [PATCH 4/5] Free iterated vec on interator exhaustion --- mypyc/lib-rt/vecs/vec_nested.c | 1 + mypyc/lib-rt/vecs/vec_t.c | 1 + mypyc/lib-rt/vecs/vec_template.c | 7 ++++++- 3 files changed, 8 insertions(+), 1 deletion(-) diff --git a/mypyc/lib-rt/vecs/vec_nested.c b/mypyc/lib-rt/vecs/vec_nested.c index 568190810012e..d7064ecbd750f 100644 --- a/mypyc/lib-rt/vecs/vec_nested.c +++ b/mypyc/lib-rt/vecs/vec_nested.c @@ -472,6 +472,7 @@ static PyObject *VecNestedIter_next(VecNestedIterObject *self) { self->index++; return item; } + Py_CLEAR(self->vec_obj); return NULL; // StopIteration } diff --git a/mypyc/lib-rt/vecs/vec_t.c b/mypyc/lib-rt/vecs/vec_t.c index c85eec42be70e..757aabf3dd673 100644 --- a/mypyc/lib-rt/vecs/vec_t.c +++ b/mypyc/lib-rt/vecs/vec_t.c @@ -467,6 +467,7 @@ static PyObject *VecTIter_next(VecTIterObject *self) { Py_INCREF(item); return item; } + Py_CLEAR(self->vec_obj); return NULL; // StopIteration } diff --git a/mypyc/lib-rt/vecs/vec_template.c b/mypyc/lib-rt/vecs/vec_template.c index c21d4af9a4934..ef2f849fefa42 100644 --- a/mypyc/lib-rt/vecs/vec_template.c +++ b/mypyc/lib-rt/vecs/vec_template.c @@ -375,11 +375,13 @@ static PyObject *vec_iter(PyObject *self) { } static void vec_iter_dealloc(NAME(IterObject) *self) { - Py_DECREF(self->vec_obj); + Py_XDECREF(self->vec_obj); PyObject_Del(self); } static PyObject *vec_iter_next(NAME(IterObject) *self) { + if (self->vec_obj == NULL) + return NULL; VEC v = self->vec_obj->vec; if (self->index < v.len) { PyObject *item = BOX_ITEM(v.buf->items[self->index]); @@ -388,10 +390,13 @@ static PyObject *vec_iter_next(NAME(IterObject) *self) { self->index++; return item; } + Py_CLEAR(self->vec_obj); return NULL; // StopIteration } static PyObject *vec_iter_len(NAME(IterObject) *self, PyObject *Py_UNUSED(ignored)) { + if (self->vec_obj == NULL) + return PyLong_FromSsize_t(0); Py_ssize_t remaining = self->vec_obj->vec.len - self->index; if (remaining < 0) remaining = 0; From 99d617e3fdb8e9031578bb317df08d9f74ec284c Mon Sep 17 00:00:00 2001 From: Jukka Lehtosalo Date: Wed, 1 Apr 2026 11:24:16 +0100 Subject: [PATCH 5/5] Store unboxed vec in iterator --- mypyc/lib-rt/vecs/vec_nested.c | 25 ++++++++++++------------- mypyc/lib-rt/vecs/vec_t.c | 25 ++++++++++++------------- mypyc/lib-rt/vecs/vec_template.c | 21 ++++++++++----------- 3 files changed, 34 insertions(+), 37 deletions(-) diff --git a/mypyc/lib-rt/vecs/vec_nested.c b/mypyc/lib-rt/vecs/vec_nested.c index d7064ecbd750f..459a40735744a 100644 --- a/mypyc/lib-rt/vecs/vec_nested.c +++ b/mypyc/lib-rt/vecs/vec_nested.c @@ -424,7 +424,7 @@ static PyMethodDef vec_methods[] = { typedef struct { PyObject_HEAD - VecNestedObject *vec_obj; // Reference to vec object (keeps buffer alive) + VecNested vec; // Unboxed vec (keeps buffer alive via buf reference) Py_ssize_t index; // Current iteration index } VecNestedIterObject; @@ -434,8 +434,8 @@ static PyObject *VecNested_iter(PyObject *self) { VecNestedIterObject *it = PyObject_GC_New(VecNestedIterObject, &VecNestedIterType); if (it == NULL) return NULL; - Py_INCREF(self); - it->vec_obj = (VecNestedObject *)self; + it->vec = ((VecNestedObject *)self)->vec; + Py_INCREF(it->vec.buf); it->index = 0; PyObject_GC_Track(it); return (PyObject *)it; @@ -444,42 +444,41 @@ static PyObject *VecNested_iter(PyObject *self) { static int VecNestedIter_traverse(VecNestedIterObject *self, visitproc visit, void *arg) { - Py_VISIT(self->vec_obj); + Py_VISIT(self->vec.buf); return 0; } static int VecNestedIter_clear(VecNestedIterObject *self) { - Py_CLEAR(self->vec_obj); + Py_CLEAR(self->vec.buf); return 0; } static void VecNestedIter_dealloc(VecNestedIterObject *self) { PyObject_GC_UnTrack(self); - Py_XDECREF(self->vec_obj); + Py_XDECREF(self->vec.buf); PyObject_GC_Del(self); } static PyObject *VecNestedIter_next(VecNestedIterObject *self) { - if (self->vec_obj == NULL) + if (self->vec.buf == NULL) return NULL; - VecNested v = self->vec_obj->vec; - if (self->index < v.len) { - PyObject *item = box_vec_item_by_index(v, self->index); + if (self->index < self->vec.len) { + PyObject *item = box_vec_item_by_index(self->vec, self->index); if (item == NULL) return NULL; self->index++; return item; } - Py_CLEAR(self->vec_obj); + Py_CLEAR(self->vec.buf); return NULL; // StopIteration } static PyObject *VecNestedIter_len(VecNestedIterObject *self, PyObject *Py_UNUSED(ignored)) { - if (self->vec_obj == NULL) + if (self->vec.buf == NULL) return PyLong_FromSsize_t(0); - Py_ssize_t remaining = self->vec_obj->vec.len - self->index; + Py_ssize_t remaining = self->vec.len - self->index; if (remaining < 0) remaining = 0; return PyLong_FromSsize_t(remaining); diff --git a/mypyc/lib-rt/vecs/vec_t.c b/mypyc/lib-rt/vecs/vec_t.c index 757aabf3dd673..79938cb9aade3 100644 --- a/mypyc/lib-rt/vecs/vec_t.c +++ b/mypyc/lib-rt/vecs/vec_t.c @@ -420,7 +420,7 @@ static PyMethodDef vec_methods[] = { typedef struct { PyObject_HEAD - VecTObject *vec_obj; // Reference to vec object (keeps buffer alive) + VecT vec; // Unboxed vec (keeps buffer alive via buf reference) Py_ssize_t index; // Current iteration index } VecTIterObject; @@ -430,8 +430,8 @@ static PyObject *VecT_iter(PyObject *self) { VecTIterObject *it = PyObject_GC_New(VecTIterObject, &VecTIterType); if (it == NULL) return NULL; - Py_INCREF(self); - it->vec_obj = (VecTObject *)self; + it->vec = ((VecTObject *)self)->vec; + Py_INCREF(it->vec.buf); it->index = 0; PyObject_GC_Track(it); return (PyObject *)it; @@ -440,41 +440,40 @@ static PyObject *VecT_iter(PyObject *self) { static int VecTIter_traverse(VecTIterObject *self, visitproc visit, void *arg) { - Py_VISIT(self->vec_obj); + Py_VISIT(self->vec.buf); return 0; } static int VecTIter_clear(VecTIterObject *self) { - Py_CLEAR(self->vec_obj); + Py_CLEAR(self->vec.buf); return 0; } static void VecTIter_dealloc(VecTIterObject *self) { PyObject_GC_UnTrack(self); - Py_XDECREF(self->vec_obj); + Py_XDECREF(self->vec.buf); PyObject_GC_Del(self); } static PyObject *VecTIter_next(VecTIterObject *self) { - if (self->vec_obj == NULL) + if (self->vec.buf == NULL) return NULL; - VecT v = self->vec_obj->vec; - if (self->index < v.len) { - PyObject *item = v.buf->items[self->index]; + if (self->index < self->vec.len) { + PyObject *item = self->vec.buf->items[self->index]; self->index++; Py_INCREF(item); return item; } - Py_CLEAR(self->vec_obj); + Py_CLEAR(self->vec.buf); return NULL; // StopIteration } static PyObject *VecTIter_len(VecTIterObject *self, PyObject *Py_UNUSED(ignored)) { - if (self->vec_obj == NULL) + if (self->vec.buf == NULL) return PyLong_FromSsize_t(0); - Py_ssize_t remaining = self->vec_obj->vec.len - self->index; + Py_ssize_t remaining = self->vec.len - self->index; if (remaining < 0) remaining = 0; return PyLong_FromSsize_t(remaining); diff --git a/mypyc/lib-rt/vecs/vec_template.c b/mypyc/lib-rt/vecs/vec_template.c index ef2f849fefa42..7881bb496a88c 100644 --- a/mypyc/lib-rt/vecs/vec_template.c +++ b/mypyc/lib-rt/vecs/vec_template.c @@ -358,7 +358,7 @@ static PyMethodDef vec_methods[] = { typedef struct { PyObject_HEAD - VEC_OBJECT *vec_obj; // Reference to vec object (keeps buffer alive) + VEC vec; // Unboxed vec (keeps buffer alive via buf reference) Py_ssize_t index; // Current iteration index } NAME(IterObject); @@ -368,36 +368,35 @@ static PyObject *vec_iter(PyObject *self) { NAME(IterObject) *it = PyObject_New(NAME(IterObject), &NAME(IterType)); if (it == NULL) return NULL; - Py_INCREF(self); - it->vec_obj = (VEC_OBJECT *)self; + it->vec = ((VEC_OBJECT *)self)->vec; + Py_XINCREF(it->vec.buf); it->index = 0; return (PyObject *)it; } static void vec_iter_dealloc(NAME(IterObject) *self) { - Py_XDECREF(self->vec_obj); + Py_XDECREF(self->vec.buf); PyObject_Del(self); } static PyObject *vec_iter_next(NAME(IterObject) *self) { - if (self->vec_obj == NULL) + if (self->vec.buf == NULL) return NULL; - VEC v = self->vec_obj->vec; - if (self->index < v.len) { - PyObject *item = BOX_ITEM(v.buf->items[self->index]); + if (self->index < self->vec.len) { + PyObject *item = BOX_ITEM(self->vec.buf->items[self->index]); if (item == NULL) return NULL; self->index++; return item; } - Py_CLEAR(self->vec_obj); + Py_CLEAR(self->vec.buf); return NULL; // StopIteration } static PyObject *vec_iter_len(NAME(IterObject) *self, PyObject *Py_UNUSED(ignored)) { - if (self->vec_obj == NULL) + if (self->vec.buf == NULL) return PyLong_FromSsize_t(0); - Py_ssize_t remaining = self->vec_obj->vec.len - self->index; + Py_ssize_t remaining = self->vec.len - self->index; if (remaining < 0) remaining = 0; return PyLong_FromSsize_t(remaining);