[mypyc] Add vec iterators (#21135)

Add specialized `librt.vecs.vec` iterator classes. There are no
primitives for `iter(v)` yet, though we could add them later. Generally
using `iter(v)` should be avoided when using vecs, since the generic
iterator protocol is inefficient for unboxed items.

I used coding agent assist but reviewed generated code carefully (this
is mostly boilerplate).

Author: JukkaL

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) {

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;

mypyc/lib-rt/vecs/vec_nested.c

@@ -420,6 +420,89 @@ static PyMethodDef vec_methods[] = {
     {NULL, NULL, 0, NULL},  /* Sentinel */
 };
 
+// Iterator type for nested vecs
+
+typedef struct {
+    PyObject_HEAD
+    VecNested vec;             // Unboxed vec (keeps buffer alive via buf reference)
+    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;
+    it->vec = ((VecNestedObject *)self)->vec;
+    Py_INCREF(it->vec.buf);
+    it->index = 0;
+    PyObject_GC_Track(it);
+    return (PyObject *)it;
+}
+
+static int
+VecNestedIter_traverse(VecNestedIterObject *self, visitproc visit, void *arg)
+{
+    Py_VISIT(self->vec.buf);
+    return 0;
+}
+
+static int
+VecNestedIter_clear(VecNestedIterObject *self)
+{
+    Py_CLEAR(self->vec.buf);
+    return 0;
+}
+
+static void VecNestedIter_dealloc(VecNestedIterObject *self) {
+    PyObject_GC_UnTrack(self);
+    Py_XDECREF(self->vec.buf);
+    PyObject_GC_Del(self);
+}
+
+static PyObject *VecNestedIter_next(VecNestedIterObject *self) {
+    if (self->vec.buf == NULL)
+        return NULL;
+    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.buf);
+    return NULL;  // StopIteration
+}
+
+static PyObject *VecNestedIter_len(VecNestedIterObject *self, PyObject *Py_UNUSED(ignored)) {
+    if (self->vec.buf == NULL)
+        return PyLong_FromSsize_t(0);
+    Py_ssize_t remaining = self->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 +530,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,

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
+    VecT vec;             // Unboxed vec (keeps buffer alive via buf reference)
+    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;
+    it->vec = ((VecTObject *)self)->vec;
+    Py_INCREF(it->vec.buf);
+    it->index = 0;
+    PyObject_GC_Track(it);
+    return (PyObject *)it;
+}
+
+static int
+VecTIter_traverse(VecTIterObject *self, visitproc visit, void *arg)
+{
+    Py_VISIT(self->vec.buf);
+    return 0;
+}
+
+static int
+VecTIter_clear(VecTIterObject *self)
+{
+    Py_CLEAR(self->vec.buf);
+    return 0;
+}
+
+static void VecTIter_dealloc(VecTIterObject *self) {
+    PyObject_GC_UnTrack(self);
+    Py_XDECREF(self->vec.buf);
+    PyObject_GC_Del(self);
+}
+
+static PyObject *VecTIter_next(VecTIterObject *self) {
+    if (self->vec.buf == NULL)
+        return NULL;
+    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.buf);
+    return NULL;  // StopIteration
+}
+
+static PyObject *VecTIter_len(VecTIterObject *self, PyObject *Py_UNUSED(ignored)) {
+    if (self->vec.buf == NULL)
+        return PyLong_FromSsize_t(0);
+    Py_ssize_t remaining = self->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,

mypyc/lib-rt/vecs/vec_template.c

@@ -354,6 +354,71 @@ static PyMethodDef vec_methods[] = {
     {NULL, NULL, 0, NULL},  /* Sentinel */
 };
 
+// Iterator type for specialized vec types
+
+typedef struct {
+    PyObject_HEAD
+    VEC vec;              // Unboxed vec (keeps buffer alive via buf reference)
+    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;
+    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.buf);
+    PyObject_Del(self);
+}
+
+static PyObject *vec_iter_next(NAME(IterObject) *self) {
+    if (self->vec.buf == NULL)
+        return NULL;
+    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.buf);
+    return NULL;  // StopIteration
+}
+
+static PyObject *vec_iter_len(NAME(IterObject) *self, PyObject *Py_UNUSED(ignored)) {
+    if (self->vec.buf == NULL)
+        return PyLong_FromSsize_t(0);
+    Py_ssize_t remaining = self->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 +442,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,

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])