diff -r 77e621214488 Modules/itertoolsmodule.c --- a/Modules/itertoolsmodule.c Sat Aug 15 15:16:12 2015 -0700 +++ b/Modules/itertoolsmodule.c Sun Aug 16 11:20:03 2015 +0300 @@ -863,7 +863,8 @@ typedef struct { PyObject_HEAD PyObject *it; PyObject *saved; - int firstpass; + Py_ssize_t index; + int legacy; } cycleobject; static PyTypeObject cycle_type; @@ -902,7 +903,8 @@ cycle_new(PyTypeObject *type, PyObject * } lz->it = it; lz->saved = saved; - lz->firstpass = 0; + lz->index = 0; + lz->legacy = 0; return (PyObject *)lz; } @@ -911,15 +913,16 @@ static void 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 @@ static PyObject * 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,81 @@ cycle_next(cycleobject *lz) } /* 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; + item = PyList_GET_ITEM(lz->saved, lz->index); + lz->index++; + if (lz->index >= Py_SIZE(lz->saved)) + lz->index = 0; + 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; + + /* create an empty iter */ + seq = PyTuple_New(0); + if (seq == NULL) + return NULL; + it = PyObject_GetIter(seq); + Py_DECREF(seq); + if (it == NULL) + return NULL; + + if (lz->index == 0) { + seq = lz->saved; + Py_INCREF(seq); + } + else { + /* make a circular shift */ + PyObject **items; + Py_ssize_t i, size = Py_SIZE(lz->saved); + seq = PyList_New(size); + if (seq == NULL) { + Py_DECREF(it); + return NULL; + } + if (Py_SIZE(lz->saved) != size) { + PyErr_SetString(PyExc_SystemError, + "list changed size during pickling"); + Py_DECREF(it); + return NULL; + } + items = PySequence_Fast_ITEMS(lz->saved); + for (i = 0; i < size; i++) + Py_INCREF(items[i]); + memcpy(PySequence_Fast_ITEMS(seq), items + lz->index, + (size - lz->index)*sizeof(PyObject*)); + memcpy(PySequence_Fast_ITEMS(seq) + (size - lz->index), items, + lz->index*sizeof(PyObject*)); + } + + return Py_BuildValue("O(N)(NO)", Py_TYPE(lz), it, seq, Py_False); + } + return Py_BuildValue("O(O)(OO)", Py_TYPE(lz), lz->it, lz->saved, + lz->legacy ? Py_True : Py_False); } 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_INCREF(saved); 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; }