diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -863,7 +863,8 @@ PyObject_HEAD PyObject *it; PyObject *saved; - int firstpass; + Py_ssize_t index; + int legacy; } cycleobject; static PyTypeObject cycle_type; @@ -902,7 +903,8 @@ } lz->it = it; lz->saved = saved; - lz->firstpass = 0; + lz->index = 0; + lz->legacy = 0; return (PyObject *)lz; } @@ -911,15 +913,16 @@ cycle_dealloc(cycleobject *lz) { PyObject_GC_UnTrack(lz); + Py_XDECREF(lz->it); Py_XDECREF(lz->saved); - Py_XDECREF(lz->it); Py_TYPE(lz)->tp_free(lz); } static int cycle_traverse(cycleobject *lz, visitproc visit, void *arg) { - Py_VISIT(lz->it); + if (lz->it) + Py_VISIT(lz->it); Py_VISIT(lz->saved); return 0; } @@ -928,13 +931,13 @@ cycle_next(cycleobject *lz) { PyObject *item; - PyObject *it; - PyObject *tmp; - - while (1) { + + if (lz->it != NULL) { item = PyIter_Next(lz->it); if (item != NULL) { - if (!lz->firstpass && PyList_Append(lz->saved, item)) { + if (lz->legacy) + return item; + if (PyList_Append(lz->saved, item)) { Py_DECREF(item); return NULL; } @@ -942,40 +945,60 @@ } /* Note: StopIteration is already cleared by PyIter_Next() */ if (PyErr_Occurred()) - return NULL; - if (PyList_Size(lz->saved) == 0) return NULL; - it = PyObject_GetIter(lz->saved); - if (it == NULL) - return NULL; - tmp = lz->it; - lz->it = it; - lz->firstpass = 1; - Py_DECREF(tmp); + Py_CLEAR(lz->it); } + if (Py_SIZE(lz->saved) == 0) + return NULL; + if (lz->index >= Py_SIZE(lz->saved)) + lz->index = 0; + item = PyList_GET_ITEM(lz->saved, lz->index); + lz->index++; + Py_INCREF(item); + return item; } static PyObject * cycle_reduce(cycleobject *lz) { - /* Create a new cycle with the iterator tuple, then set - * the saved state on it. - */ - return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz), - lz->it, lz->saved, lz->firstpass); + /* Create a new cycle with the iterator tuple, then set the saved state */ + if (lz->it == NULL) { + PyObject *seq, *it, *result; + + if (lz->index >= Py_SIZE(lz->saved)) + lz->index = 0; + seq = PyList_GetSlice(lz->saved, lz->index, Py_SIZE(lz->saved)); + if (seq == NULL) + return NULL; + it = PyObject_GetIter(seq); + Py_DECREF(seq); + if (it == NULL) + return NULL; + seq = PyList_GetSlice(lz->saved, 0, lz->index); + if (seq == NULL) { + Py_DECREF(it); + return NULL; + } + result = Py_BuildValue("O(O)(Oi)", Py_TYPE(lz), it, seq, 0); + Py_DECREF(it); + Py_DECREF(seq); + return result; + } + return Py_BuildValue("O(O)(Oi)", Py_TYPE(lz), lz->it, lz->saved, 0); } static PyObject * cycle_setstate(cycleobject *lz, PyObject *state) { PyObject *saved=NULL; - int firstpass; - if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &firstpass)) + int legacy; + if (!PyArg_ParseTuple(state, "O!i", &PyList_Type, &saved, &legacy)) return NULL; Py_CLEAR(lz->saved); lz->saved = saved; Py_XINCREF(lz->saved); - lz->firstpass = firstpass != 0; + lz->legacy = legacy != 0; + lz->index = 0; Py_RETURN_NONE; }