gh-118218: Reuse return tuple in itertools.pairwise (GH-118219)

diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index 95e67911db6a7f2df3dc40af42caa760ee35cea1..e243da309f03d8b6264a17491c42a2badcfa3d60 100644 (file)
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -1821,6 +1821,13 @@ class TestBasicOps(unittest.TestCase):
         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

diff --git a/Misc/NEWS.d/next/Library/2024-04-24-07-45-08.gh-issue-118218.m1OHbN.rst b/Misc/NEWS.d/next/Library/2024-04-24-07-45-08.gh-issue-118218.m1OHbN.rst
new file mode 100644 (file)
index 0000000..6d3c54c
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2024-04-24-07-45-08.gh-issue-118218.m1OHbN.rst
@@ -0,0 +1 @@
+Speed up :func:`itertools.pairwise` in the common case by up to 1.8x.

diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c
index 6ee447ef6a8cd6c4cfe336dd52c240cbeb62e432..21ce3ecfad0354744d2d33af8cd1490999aacd49 100644 (file)
--- a/Modules/itertoolsmodule.c
+++ b/Modules/itertoolsmodule.c
@@ -270,6 +270,7 @@ typedef struct {
     PyObject_HEAD
     PyObject *it;
     PyObject *old;
+    PyObject *result;
 } pairwiseobject;
 
 /*[clinic input]
@@ -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;
 }
 
@@ -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);
 }
@@ -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;
 }
 
@@ -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;