Skip to content

Commit

Permalink
gh-118218: Reuse return tuple in itertools.pairwise (GH-118219)
Browse files Browse the repository at this point in the history
  • Loading branch information
hauntsaninja committed Apr 30, 2024
1 parent b568c2c commit 6999d68
Show file tree
Hide file tree
Showing 3 changed files with 40 additions and 2 deletions.
7 changes: 7 additions & 0 deletions Lib/test/test_itertools.py
Original file line number Diff line number Diff line change
Expand Up @@ -1821,6 +1821,13 @@ def test_zip_longest_result_gc(self):
gc.collect()
self.assertTrue(gc.is_tracked(next(it)))

@support.cpython_only
def test_pairwise_result_gc(self):
# Ditto for pairwise.
it = pairwise([None, None])
gc.collect()
self.assertTrue(gc.is_tracked(next(it)))

@support.cpython_only
def test_immutable_types(self):
from itertools import _grouper, _tee, _tee_dataobject
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Speed up :func:`itertools.pairwise` in the common case by up to 1.8x.
34 changes: 32 additions & 2 deletions Modules/itertoolsmodule.c
Original file line number Diff line number Diff line change
Expand Up @@ -270,6 +270,7 @@ typedef struct {
PyObject_HEAD
PyObject *it;
PyObject *old;
PyObject *result;

This comment has been minimized.

Copy link
@pochmann3

pochmann3 May 1, 2024

Contributor

Hmm, now the old object is stored twice, right? Both in old and in result. Is that intentional?

} pairwiseobject;

/*[clinic input]
Expand Down Expand Up @@ -301,6 +302,11 @@ pairwise_new_impl(PyTypeObject *type, PyObject *iterable)
}
po->it = it;
po->old = NULL;
po->result = PyTuple_Pack(2, Py_None, Py_None);
if (po->result == NULL) {
Py_DECREF(po);
return NULL;
}
return (PyObject *)po;
}

Expand All @@ -311,6 +317,7 @@ pairwise_dealloc(pairwiseobject *po)
PyObject_GC_UnTrack(po);
Py_XDECREF(po->it);
Py_XDECREF(po->old);
Py_XDECREF(po->result);
tp->tp_free(po);
Py_DECREF(tp);
}
Expand All @@ -321,6 +328,7 @@ pairwise_traverse(pairwiseobject *po, visitproc visit, void *arg)
Py_VISIT(Py_TYPE(po));
Py_VISIT(po->it);
Py_VISIT(po->old);
Py_VISIT(po->result);
return 0;
}

Expand Down Expand Up @@ -355,8 +363,30 @@ pairwise_next(pairwiseobject *po)
Py_DECREF(old);
return NULL;
}
/* Future optimization: Reuse the result tuple as we do in enumerate() */
result = PyTuple_Pack(2, old, new);

result = po->result;
if (Py_REFCNT(result) == 1) {
Py_INCREF(result);
PyObject *last_old = PyTuple_GET_ITEM(result, 0);
PyObject *last_new = PyTuple_GET_ITEM(result, 1);
PyTuple_SET_ITEM(result, 0, Py_NewRef(old));
PyTuple_SET_ITEM(result, 1, Py_NewRef(new));
Py_DECREF(last_old);
Py_DECREF(last_new);
// bpo-42536: The GC may have untracked this result tuple. Since we're
// recycling it, make sure it's tracked again:
if (!_PyObject_GC_IS_TRACKED(result)) {
_PyObject_GC_TRACK(result);
}
}
else {
result = PyTuple_New(2);
if (result != NULL) {
PyTuple_SET_ITEM(result, 0, Py_NewRef(old));
PyTuple_SET_ITEM(result, 1, Py_NewRef(new));
}
}

Py_XSETREF(po->old, new);
Py_DECREF(old);
return result;
Expand Down

0 comments on commit 6999d68

Please sign in to comment.