Skip to content

Commit 3fbbd55

Browse files
JukkaLJukkaL
authored
[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).
1 parent a25c331 commit 3fbbd55

File tree

12 files changed

+557
-0
lines changed

12 files changed

+557
-0
lines changed

‎mypyc/lib-rt/vecs/librt_vecs.c‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -950,36 +950,52 @@ librt_vecs_module_exec(PyObject *m)
950950
return -1;
951951
if (PyType_Ready(&VecTBufType) < 0)
952952
return -1;
953+
if (PyType_Ready(&VecTIterType) < 0)
954+
return -1;
953955

954956
if (PyType_Ready(&VecNestedType) < 0)
955957
return -1;
956958
if (PyType_Ready(&VecNestedBufType) < 0)
957959
return -1;
960+
if (PyType_Ready(&VecNestedIterType) < 0)
961+
return -1;
958962

959963
if (PyType_Ready(&VecI64Type) < 0)
960964
return -1;
961965
if (PyType_Ready(&VecI64BufType) < 0)
962966
return -1;
967+
if (PyType_Ready(&VecI64IterType) < 0)
968+
return -1;
963969
if (PyType_Ready(&VecI32Type) < 0)
964970
return -1;
965971
if (PyType_Ready(&VecI32BufType) < 0)
966972
return -1;
973+
if (PyType_Ready(&VecI32IterType) < 0)
974+
return -1;
967975
if (PyType_Ready(&VecI16Type) < 0)
968976
return -1;
969977
if (PyType_Ready(&VecI16BufType) < 0)
970978
return -1;
979+
if (PyType_Ready(&VecI16IterType) < 0)
980+
return -1;
971981
if (PyType_Ready(&VecU8Type) < 0)
972982
return -1;
973983
if (PyType_Ready(&VecU8BufType) < 0)
974984
return -1;
985+
if (PyType_Ready(&VecU8IterType) < 0)
986+
return -1;
975987
if (PyType_Ready(&VecFloatType) < 0)
976988
return -1;
977989
if (PyType_Ready(&VecFloatBufType) < 0)
978990
return -1;
991+
if (PyType_Ready(&VecFloatIterType) < 0)
992+
return -1;
979993
if (PyType_Ready(&VecBoolType) < 0)
980994
return -1;
981995
if (PyType_Ready(&VecBoolBufType) < 0)
982996
return -1;
997+
if (PyType_Ready(&VecBoolIterType) < 0)
998+
return -1;
983999

9841000
Py_INCREF(&VecType);
9851001
if (PyModule_AddObject(m, "vec", (PyObject *)&VecType) < 0) {

‎mypyc/lib-rt/vecs/librt_vecs.h‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -493,6 +493,16 @@ extern PyTypeObject VecBoolType;
493493
extern PyTypeObject VecTType;
494494
extern PyTypeObject VecNestedType;
495495

496+
// Iterator type objects for vec iteration
497+
extern PyTypeObject VecI64IterType;
498+
extern PyTypeObject VecI32IterType;
499+
extern PyTypeObject VecI16IterType;
500+
extern PyTypeObject VecU8IterType;
501+
extern PyTypeObject VecFloatIterType;
502+
extern PyTypeObject VecBoolIterType;
503+
extern PyTypeObject VecTIterType;
504+
extern PyTypeObject VecNestedIterType;
505+
496506
// Type objects corresponding to the 'i64', 'i32', 'i16, and 'u8' types
497507
extern PyTypeObject *LibRTVecs_I64TypeObj;
498508
extern PyTypeObject *LibRTVecs_I32TypeObj;

‎mypyc/lib-rt/vecs/vec_nested.c‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -420,6 +420,89 @@ static PyMethodDef vec_methods[] = {
420420
{NULL, NULL, 0, NULL}, /* Sentinel */
421421
};
422422

423+
// Iterator type for nested vecs
424+
425+
typedef struct {
426+
PyObject_HEAD
427+
VecNested vec; // Unboxed vec (keeps buffer alive via buf reference)
428+
Py_ssize_t index; // Current iteration index
429+
} VecNestedIterObject;
430+
431+
PyTypeObject VecNestedIterType;
432+
433+
static PyObject *VecNested_iter(PyObject *self) {
434+
VecNestedIterObject *it = PyObject_GC_New(VecNestedIterObject, &VecNestedIterType);
435+
if (it == NULL)
436+
return NULL;
437+
it->vec = ((VecNestedObject *)self)->vec;
438+
Py_INCREF(it->vec.buf);
439+
it->index = 0;
440+
PyObject_GC_Track(it);
441+
return (PyObject *)it;
442+
}
443+
444+
static int
445+
VecNestedIter_traverse(VecNestedIterObject *self, visitproc visit, void *arg)
446+
{
447+
Py_VISIT(self->vec.buf);
448+
return 0;
449+
}
450+
451+
static int
452+
VecNestedIter_clear(VecNestedIterObject *self)
453+
{
454+
Py_CLEAR(self->vec.buf);
455+
return 0;
456+
}
457+
458+
static void VecNestedIter_dealloc(VecNestedIterObject *self) {
459+
PyObject_GC_UnTrack(self);
460+
Py_XDECREF(self->vec.buf);
461+
PyObject_GC_Del(self);
462+
}
463+
464+
static PyObject *VecNestedIter_next(VecNestedIterObject *self) {
465+
if (self->vec.buf == NULL)
466+
return NULL;
467+
if (self->index < self->vec.len) {
468+
PyObject *item = box_vec_item_by_index(self->vec, self->index);
469+
if (item == NULL)
470+
return NULL;
471+
self->index++;
472+
return item;
473+
}
474+
Py_CLEAR(self->vec.buf);
475+
return NULL; // StopIteration
476+
}
477+
478+
static PyObject *VecNestedIter_len(VecNestedIterObject *self, PyObject *Py_UNUSED(ignored)) {
479+
if (self->vec.buf == NULL)
480+
return PyLong_FromSsize_t(0);
481+
Py_ssize_t remaining = self->vec.len - self->index;
482+
if (remaining < 0)
483+
remaining = 0;
484+
return PyLong_FromSsize_t(remaining);
485+
}
486+
487+
static PyMethodDef VecNestedIter_methods[] = {
488+
{"__length_hint__", (PyCFunction)VecNestedIter_len, METH_NOARGS, NULL},
489+
{NULL, NULL, 0, NULL},
490+
};
491+
492+
PyTypeObject VecNestedIterType = {
493+
PyVarObject_HEAD_INIT(NULL, 0)
494+
.tp_name = "vec_nested_iterator",
495+
.tp_basicsize = sizeof(VecNestedIterObject),
496+
.tp_itemsize = 0,
497+
.tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
498+
.tp_traverse = (traverseproc)VecNestedIter_traverse,
499+
.tp_clear = (inquiry)VecNestedIter_clear,
500+
.tp_dealloc = (destructor)VecNestedIter_dealloc,
501+
.tp_iter = PyObject_SelfIter,
502+
.tp_iternext = (iternextfunc)VecNestedIter_next,
503+
.tp_methods = VecNestedIter_methods,
504+
};
505+
423506
PyTypeObject VecNestedBufType = {
424507
PyVarObject_HEAD_INIT(NULL, 0)
425508
.tp_name = "vecbuf",
@@ -447,6 +530,7 @@ PyTypeObject VecNestedType = {
447530
.tp_dealloc = (destructor)VecNested_dealloc,
448531
//.tp_free = PyObject_GC_Del,
449532
.tp_repr = (reprfunc)vec_repr,
533+
.tp_iter = VecNested_iter,
450534
.tp_as_sequence = &VecNestedSequence,
451535
.tp_as_mapping = &VecNestedMapping,
452536
.tp_richcompare = vec_richcompare,

‎mypyc/lib-rt/vecs/vec_t.c‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -416,6 +416,88 @@ static PyMethodDef vec_methods[] = {
416416
{NULL, NULL, 0, NULL}, /* Sentinel */
417417
};
418418

419+
// Iterator type for vec[T] (reference types)
420+
421+
typedef struct {
422+
PyObject_HEAD
423+
VecT vec; // Unboxed vec (keeps buffer alive via buf reference)
424+
Py_ssize_t index; // Current iteration index
425+
} VecTIterObject;
426+
427+
PyTypeObject VecTIterType;
428+
429+
static PyObject *VecT_iter(PyObject *self) {
430+
VecTIterObject *it = PyObject_GC_New(VecTIterObject, &VecTIterType);
431+
if (it == NULL)
432+
return NULL;
433+
it->vec = ((VecTObject *)self)->vec;
434+
Py_INCREF(it->vec.buf);
435+
it->index = 0;
436+
PyObject_GC_Track(it);
437+
return (PyObject *)it;
438+
}
439+
440+
static int
441+
VecTIter_traverse(VecTIterObject *self, visitproc visit, void *arg)
442+
{
443+
Py_VISIT(self->vec.buf);
444+
return 0;
445+
}
446+
447+
static int
448+
VecTIter_clear(VecTIterObject *self)
449+
{
450+
Py_CLEAR(self->vec.buf);
451+
return 0;
452+
}
453+
454+
static void VecTIter_dealloc(VecTIterObject *self) {
455+
PyObject_GC_UnTrack(self);
456+
Py_XDECREF(self->vec.buf);
457+
PyObject_GC_Del(self);
458+
}
459+
460+
static PyObject *VecTIter_next(VecTIterObject *self) {
461+
if (self->vec.buf == NULL)
462+
return NULL;
463+
if (self->index < self->vec.len) {
464+
PyObject *item = self->vec.buf->items[self->index];
465+
self->index++;
466+
Py_INCREF(item);
467+
return item;
468+
}
469+
Py_CLEAR(self->vec.buf);
470+
return NULL; // StopIteration
471+
}
472+
473+
static PyObject *VecTIter_len(VecTIterObject *self, PyObject *Py_UNUSED(ignored)) {
474+
if (self->vec.buf == NULL)
475+
return PyLong_FromSsize_t(0);
476+
Py_ssize_t remaining = self->vec.len - self->index;
477+
if (remaining < 0)
478+
remaining = 0;
479+
return PyLong_FromSsize_t(remaining);
480+
}
481+
482+
static PyMethodDef VecTIter_methods[] = {
483+
{"__length_hint__", (PyCFunction)VecTIter_len, METH_NOARGS, NULL},
484+
{NULL, NULL, 0, NULL},
485+
};
486+
487+
PyTypeObject VecTIterType = {
488+
PyVarObject_HEAD_INIT(NULL, 0)
489+
.tp_name = "vec_iterator",
490+
.tp_basicsize = sizeof(VecTIterObject),
491+
.tp_itemsize = 0,
492+
.tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
493+
.tp_traverse = (traverseproc)VecTIter_traverse,
494+
.tp_clear = (inquiry)VecTIter_clear,
495+
.tp_dealloc = (destructor)VecTIter_dealloc,
496+
.tp_iter = PyObject_SelfIter,
497+
.tp_iternext = (iternextfunc)VecTIter_next,
498+
.tp_methods = VecTIter_methods,
499+
};
500+
419501
PyTypeObject VecTBufType = {
420502
PyVarObject_HEAD_INIT(NULL, 0)
421503
.tp_name = "vecbuf",
@@ -443,6 +525,7 @@ PyTypeObject VecTType = {
443525
.tp_dealloc = (destructor)VecT_dealloc,
444526
//.tp_free = PyObject_GC_Del,
445527
.tp_repr = (reprfunc)vec_repr,
528+
.tp_iter = VecT_iter,
446529
.tp_as_sequence = &VecTSequence,
447530
.tp_as_mapping = &VecTMapping,
448531
.tp_richcompare = vec_richcompare,

‎mypyc/lib-rt/vecs/vec_template.c‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -354,6 +354,71 @@ static PyMethodDef vec_methods[] = {
354354
{NULL, NULL, 0, NULL}, /* Sentinel */
355355
};
356356

357+
// Iterator type for specialized vec types
358+
359+
typedef struct {
360+
PyObject_HEAD
361+
VEC vec; // Unboxed vec (keeps buffer alive via buf reference)
362+
Py_ssize_t index; // Current iteration index
363+
} NAME(IterObject);
364+
365+
PyTypeObject NAME(IterType);
366+
367+
static PyObject *vec_iter(PyObject *self) {
368+
NAME(IterObject) *it = PyObject_New(NAME(IterObject), &NAME(IterType));
369+
if (it == NULL)
370+
return NULL;
371+
it->vec = ((VEC_OBJECT *)self)->vec;
372+
Py_XINCREF(it->vec.buf);
373+
it->index = 0;
374+
return (PyObject *)it;
375+
}
376+
377+
static void vec_iter_dealloc(NAME(IterObject) *self) {
378+
Py_XDECREF(self->vec.buf);
379+
PyObject_Del(self);
380+
}
381+
382+
static PyObject *vec_iter_next(NAME(IterObject) *self) {
383+
if (self->vec.buf == NULL)
384+
return NULL;
385+
if (self->index < self->vec.len) {
386+
PyObject *item = BOX_ITEM(self->vec.buf->items[self->index]);
387+
if (item == NULL)
388+
return NULL;
389+
self->index++;
390+
return item;
391+
}
392+
Py_CLEAR(self->vec.buf);
393+
return NULL; // StopIteration
394+
}
395+
396+
static PyObject *vec_iter_len(NAME(IterObject) *self, PyObject *Py_UNUSED(ignored)) {
397+
if (self->vec.buf == NULL)
398+
return PyLong_FromSsize_t(0);
399+
Py_ssize_t remaining = self->vec.len - self->index;
400+
if (remaining < 0)
401+
remaining = 0;
402+
return PyLong_FromSsize_t(remaining);
403+
}
404+
405+
static PyMethodDef vec_iter_methods[] = {
406+
{"__length_hint__", (PyCFunction)vec_iter_len, METH_NOARGS, NULL},
407+
{NULL, NULL, 0, NULL},
408+
};
409+
410+
PyTypeObject NAME(IterType) = {
411+
PyVarObject_HEAD_INIT(NULL, 0)
412+
.tp_name = "vec_iterator[" ITEM_TYPE_STR "]",
413+
.tp_basicsize = sizeof(NAME(IterObject)),
414+
.tp_itemsize = 0,
415+
.tp_flags = Py_TPFLAGS_DEFAULT,
416+
.tp_dealloc = (destructor)vec_iter_dealloc,
417+
.tp_iter = PyObject_SelfIter,
418+
.tp_iternext = (iternextfunc)vec_iter_next,
419+
.tp_methods = vec_iter_methods,
420+
};
421+
357422
PyTypeObject BUF_TYPE = {
358423
PyVarObject_HEAD_INIT(NULL, 0)
359424
.tp_name = "vecbuf[" ITEM_TYPE_STR "]",
@@ -377,6 +442,7 @@ PyTypeObject VEC_TYPE = {
377442
//.tp_free = PyObject_Del,
378443
.tp_dealloc = (destructor)vec_dealloc,
379444
.tp_repr = (reprfunc)vec_repr,
445+
.tp_iter = vec_iter,
380446
.tp_as_sequence = &vec_sequence_methods,
381447
.tp_as_mapping = &vec_mapping_methods,
382448
.tp_richcompare = vec_richcompare,

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

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -231,6 +231,45 @@ def test_iteration_nested() -> None:
231231
(5, 2), (5, 5), (5, 7),
232232
(7, 2), (7, 5), (7, 7)]
233233

234+
def test_iterator_protocol() -> None:
235+
# Test iter() and next() functions explicitly
236+
v = vec[i64]([1, 2, 3])
237+
it = iter(v)
238+
assert next(it) == 1
239+
assert next(it) == 2
240+
assert next(it) == 3
241+
with assertRaises(StopIteration):
242+
next(it)
243+
244+
# Test multiple iterators on same vec
245+
v = vec[i64]([10, 20])
246+
it1 = iter(v)
247+
it2 = iter(v)
248+
assert next(it1) == 10
249+
assert next(it2) == 10
250+
assert next(it1) == 20
251+
assert next(it2) == 20
252+
253+
# Test empty vec iterator
254+
it = iter(vec[i64]())
255+
with assertRaises(StopIteration):
256+
next(it)
257+
258+
def test_iterator_length_hint() -> None:
259+
v = vec[i64]([1, 2, 3, 4, 5])
260+
it = iter(v)
261+
op: Any = getattr(sys, "modules")["operator"]
262+
length_hint = op.length_hint
263+
assert length_hint(it) == 5
264+
next(it)
265+
assert length_hint(it) == 4
266+
next(it)
267+
next(it)
268+
assert length_hint(it) == 2
269+
next(it)
270+
next(it)
271+
assert length_hint(it) == 0
272+
234273
def test_slicing() -> None:
235274
v = vec[i64](range(5))
236275
assert v[1:4] == vec[i64]([1, 2, 3])

0 commit comments

Comments
 (0)