gh-115999: Add free-threaded specialization for FOR_ITER

Include/internal/pycore_list.h

@@ -13,6 +13,10 @@ extern void _PyList_DebugMallocStats(FILE *out);
 // _PyList_GetItemRef should be used only when the object is known as a list
 // because it doesn't raise TypeError when the object is not a list, whereas PyList_GetItemRef does.
 extern PyObject* _PyList_GetItemRef(PyListObject *, Py_ssize_t i);
+#ifdef Py_GIL_DISABLED
+// Returns -1 in case of races with other threads.
+extern int _PyList_GetItemRefNoLock(PyListObject *, Py_ssize_t i, PyObject ** result);
+#endif
 
 #define _PyList_ITEMS(op) _Py_RVALUE(_PyList_CAST(op)->ob_item)
 

Lib/test/libregrtest/tsan.py

@@ -26,6 +26,7 @@
     'test_threadsignals',
     'test_weakref',
     'test_free_threading.test_slots',
+    'test_free_threading.test_iteration',
 ]
 
 

Objects/listobject.c

@@ -357,7 +357,7 @@ list_get_item_ref(PyListObject *op, Py_ssize_t i)
         return NULL;
     }
     Py_ssize_t cap = list_capacity(ob_item);
-    assert(cap != -1 && cap >= size);
+    assert(cap != -1);
     if (!valid_index(i, cap)) {
         return NULL;
     }
@@ -415,6 +415,32 @@ _PyList_GetItemRef(PyListObject *list, Py_ssize_t i)
     return list_get_item_ref(list, i);
 }
 
+#ifdef Py_GIL_DISABLED
+int
+_PyList_GetItemRefNoLock(PyListObject *list, Py_ssize_t i, PyObject **result)
+{
+    assert(_Py_IsOwnedByCurrentThread((PyObject *)list) ||
+           _PyObject_GC_IS_SHARED(list));
+    if (!valid_index(i, PyList_GET_SIZE(list))) {
+        return 0;
+    }
+    PyObject **ob_item = _Py_atomic_load_ptr(&list->ob_item);
+    if (ob_item == NULL) {
+        return 0;
+    }
+    Py_ssize_t cap = list_capacity(ob_item);
+    assert(cap != -1);
+    if (!valid_index(i, cap)) {
+        return 0;
+    }
+    *result = _Py_TryXGetRef(&ob_item[i]);
+    if (*result == NULL) {
+        return -1;
+    }
+    return 1;
+}
+#endif
+
 int
 PyList_SetItem(PyObject *op, Py_ssize_t i,
                PyObject *newitem)

Objects/rangeobject.c

@@ -7,6 +7,7 @@
 #include "pycore_modsupport.h"    // _PyArg_NoKwnames()
 #include "pycore_range.h"
 #include "pycore_tuple.h"         // _PyTuple_ITEMS()
+#include "pycore_pyatomic_ft_wrappers.h"
 
 
 /* Support objects whose length is > PY_SSIZE_T_MAX.
@@ -816,10 +817,12 @@ PyTypeObject PyRange_Type = {
 static PyObject *
 rangeiter_next(_PyRangeIterObject *r)
 {
-    if (r->len > 0) {
-        long result = r->start;
-        r->start = result + r->step;
-        r->len--;
+    long len = FT_ATOMIC_LOAD_LONG_RELAXED(r->len);
+    if (len > 0) {
+        long result = FT_ATOMIC_LOAD_LONG_RELAXED(r->start);
+        FT_ATOMIC_STORE_LONG_RELAXED(r->start, result + r->step);
+        // Relaxed ops for maximum speed and minimum thread-safety.
+        FT_ATOMIC_STORE_LONG_RELAXED(r->len, len - 1);
         return PyLong_FromLong(result);
     }
     return NULL;
@@ -828,7 +831,7 @@ rangeiter_next(_PyRangeIterObject *r)
 static PyObject *
 rangeiter_len(_PyRangeIterObject *r, PyObject *Py_UNUSED(ignored))
 {
-    return PyLong_FromLong(r->len);
+    return PyLong_FromLong(FT_ATOMIC_LOAD_LONG_RELAXED(r->len));
 }
 
 PyDoc_STRVAR(length_hint_doc,
@@ -841,10 +844,11 @@ rangeiter_reduce(_PyRangeIterObject *r, PyObject *Py_UNUSED(ignored))
     PyObject *range;
 
     /* create a range object for pickling */
-    start = PyLong_FromLong(r->start);
+    long lstart = FT_ATOMIC_LOAD_LONG_RELAXED(r->start);
+    start = PyLong_FromLong(lstart);
     if (start == NULL)
         goto err;
-    stop = PyLong_FromLong(r->start + r->len * r->step);
+    stop = PyLong_FromLong(lstart + FT_ATOMIC_LOAD_LONG_RELAXED(r->len) * r->step);
     if (stop == NULL)
         goto err;
     step = PyLong_FromLong(r->step);
@@ -871,12 +875,14 @@ rangeiter_setstate(_PyRangeIterObject *r, PyObject *state)
     if (index == -1 && PyErr_Occurred())
         return NULL;
     /* silently clip the index value */
+    long len = FT_ATOMIC_LOAD_LONG_RELAXED(r->len);
     if (index < 0)
         index = 0;
-    else if (index > r->len)
-        index = r->len; /* exhausted iterator */
-    r->start += index * r->step;
-    r->len -= index;
+    else if (index > len)
+        index = len; /* exhausted iterator */
+    FT_ATOMIC_STORE_LONG_RELAXED(r->start,
+        FT_ATOMIC_LOAD_LONG_RELAXED(r->start) + index * r->step);
+    FT_ATOMIC_STORE_LONG_RELAXED(r->len, len - index);
     Py_RETURN_NONE;
 }
 

Objects/tupleobject.c

@@ -1020,13 +1020,15 @@ tupleiter_next(PyObject *self)
         return NULL;
     assert(PyTuple_Check(seq));
 
-    if (it->it_index < PyTuple_GET_SIZE(seq)) {
-        item = PyTuple_GET_ITEM(seq, it->it_index);
-        ++it->it_index;
+    Py_ssize_t idx = FT_ATOMIC_LOAD_SSIZE_RELAXED(it->it_index);
+    if ((size_t)idx < (size_t)PyTuple_GET_SIZE(seq)) {
+        item = PyTuple_GET_ITEM(seq, idx++);
+        FT_ATOMIC_STORE_SSIZE_RELAXED(it->it_index, idx);
         return Py_NewRef(item);
     }
 
-    it->it_seq = NULL;
+    FT_ATOMIC_STORE_SSIZE_RELAXED(it->it_index, -1);
+    FT_ATOMIC_STORE_PTR_RELAXED(it->it_seq, NULL);
     Py_DECREF(seq);
     return NULL;
 }

Python/bytecodes.c

@@ -2929,15 +2929,15 @@ dummy_func(
         };
 
         specializing op(_SPECIALIZE_FOR_ITER, (counter/1, iter -- iter)) {
-            #if ENABLE_SPECIALIZATION
+            #if ENABLE_SPECIALIZATION_FT
             if (ADAPTIVE_COUNTER_TRIGGERS(counter)) {
                 next_instr = this_instr;
                 _Py_Specialize_ForIter(iter, next_instr, oparg);
                 DISPATCH_SAME_OPARG();
             }
             OPCODE_DEFERRED_INC(FOR_ITER);
             ADVANCE_ADAPTIVE_COUNTER(this_instr[1].counter);
-            #endif  /* ENABLE_SPECIALIZATION */
+            #endif  /* ENABLE_SPECIALIZATION_FT */
         }
 
         replaced op(_FOR_ITER, (iter -- iter, next)) {
@@ -3015,10 +3015,20 @@ dummy_func(
 
 
         op(_ITER_CHECK_LIST, (iter -- iter)) {
-            EXIT_IF(Py_TYPE(PyStackRef_AsPyObjectBorrow(iter)) != &PyListIter_Type);
+            PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
+            EXIT_IF(Py_TYPE(iter_o) != &PyListIter_Type);
+#ifdef Py_GIL_DISABLED
+            _PyListIterObject *it = (_PyListIterObject *)iter_o;
+            EXIT_IF(it->it_seq == NULL ||
+                    !_Py_IsOwnedByCurrentThread((PyObject *)it->it_seq) ||
+                    !_PyObject_GC_IS_SHARED(it->it_seq));
+#endif
         }
 
         replaced op(_ITER_JUMP_LIST, (iter -- iter)) {
+// For free-threaded Python, the loop exit can happen at any point during item
+// retrieval, so separate ops don't make much sense.
+#ifndef Py_GIL_DISABLED
             PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
             _PyListIterObject *it = (_PyListIterObject *)iter_o;
             assert(Py_TYPE(iter_o) == &PyListIter_Type);
@@ -3036,10 +3046,12 @@ dummy_func(
                 JUMPBY(oparg + 1);
                 DISPATCH();
             }
+#endif
         }
 
         // Only used by Tier 2
         op(_GUARD_NOT_EXHAUSTED_LIST, (iter -- iter)) {
+#ifndef Py_GIL_DISABLED
             PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
             _PyListIterObject *it = (_PyListIterObject *)iter_o;
             assert(Py_TYPE(iter_o) == &PyListIter_Type);
@@ -3049,6 +3061,7 @@ dummy_func(
                 it->it_index = -1;
                 EXIT_IF(1);
             }
+#endif
         }
 
         op(_ITER_NEXT_LIST, (iter -- iter, next)) {
@@ -3057,8 +3070,30 @@ dummy_func(
             assert(Py_TYPE(iter_o) == &PyListIter_Type);
             PyListObject *seq = it->it_seq;
             assert(seq);
+#ifdef Py_GIL_DISABLED
+            assert(_Py_IsOwnedByCurrentThread((PyObject *)seq) ||
+                   _PyObject_GC_IS_SHARED(seq));
+            STAT_INC(FOR_ITER, hit);
+            Py_ssize_t idx = _Py_atomic_load_ssize_relaxed(&it->it_index);
+            PyObject *item;
+            int result = _PyList_GetItemRefNoLock(it->it_seq, idx, &item);
+            // A negative result means we lost a race with another thread
+            // and we need to take the slow path.
+            EXIT_IF(result < 0);
+            if (result == 0) {
+                _Py_atomic_store_ssize_relaxed(&it->it_index, -1);
+                PyStackRef_CLOSE(iter);
+                STACK_SHRINK(1);
+                /* Jump forward oparg, then skip following END_FOR and POP_TOP instructions */
+                JUMPBY(oparg + 2);
+                DISPATCH();
+            }
+            _Py_atomic_store_ssize_relaxed(&it->it_index, idx + 1);
+            next = PyStackRef_FromPyObjectSteal(item);
+#else
             assert(it->it_index < PyList_GET_SIZE(seq));
             next = PyStackRef_FromPyObjectNew(PyList_GET_ITEM(seq, it->it_index++));
+#endif
         }
 
         macro(FOR_ITER_LIST) =
@@ -3073,8 +3108,11 @@ dummy_func(
 
         replaced op(_ITER_JUMP_TUPLE, (iter -- iter)) {
             PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
-            _PyTupleIterObject *it = (_PyTupleIterObject *)iter_o;
             assert(Py_TYPE(iter_o) == &PyTupleIter_Type);
+// For free-threaded Python, the loop exit can happen at any point during item
+// retrieval, so separate ops don't make much sense.
+#ifndef Py_GIL_DISABLED
+            _PyTupleIterObject *it = (_PyTupleIterObject *)iter_o;
             STAT_INC(FOR_ITER, hit);
             PyTupleObject *seq = it->it_seq;
             if (seq == NULL || it->it_index >= PyTuple_GET_SIZE(seq)) {
@@ -3086,26 +3124,44 @@ dummy_func(
                 JUMPBY(oparg + 1);
                 DISPATCH();
             }
+#endif
         }
 
         // Only used by Tier 2
         op(_GUARD_NOT_EXHAUSTED_TUPLE, (iter -- iter)) {
             PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
             _PyTupleIterObject *it = (_PyTupleIterObject *)iter_o;
             assert(Py_TYPE(iter_o) == &PyTupleIter_Type);
+#ifndef Py_GIL_DISABLED
             PyTupleObject *seq = it->it_seq;
             EXIT_IF(seq == NULL);
             EXIT_IF(it->it_index >= PyTuple_GET_SIZE(seq));
+#endif
         }
 
         op(_ITER_NEXT_TUPLE, (iter -- iter, next)) {
             PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
             _PyTupleIterObject *it = (_PyTupleIterObject *)iter_o;
             assert(Py_TYPE(iter_o) == &PyTupleIter_Type);
             PyTupleObject *seq = it->it_seq;
+#ifdef Py_GIL_DISABLED
+            STAT_INC(FOR_ITER, hit);
+            Py_ssize_t idx = _Py_atomic_load_ssize_relaxed(&it->it_index);
+            if (seq == NULL || (size_t)idx >= (size_t)PyTuple_GET_SIZE(seq)) {
+                _Py_atomic_store_ssize_relaxed(&it->it_index, -1);
+                PyStackRef_CLOSE(iter);
+                STACK_SHRINK(1);
+                /* Jump forward oparg, then skip following END_FOR and POP_TOP instructions */
+                JUMPBY(oparg + 2);
+                DISPATCH();
+            }
+            _Py_atomic_store_ssize_relaxed(&it->it_index, idx + 1);
+            next = PyStackRef_FromPyObjectNew(PyTuple_GET_ITEM(seq, idx));
+#else
             assert(seq);
             assert(it->it_index < PyTuple_GET_SIZE(seq));
             next = PyStackRef_FromPyObjectNew(PyTuple_GET_ITEM(seq, it->it_index++));
+#endif
         }
 
         macro(FOR_ITER_TUPLE) =
@@ -3123,7 +3179,7 @@ dummy_func(
             _PyRangeIterObject *r = (_PyRangeIterObject *)PyStackRef_AsPyObjectBorrow(iter);
             assert(Py_TYPE(r) == &PyRangeIter_Type);
             STAT_INC(FOR_ITER, hit);
-            if (r->len <= 0) {
+            if (FT_ATOMIC_LOAD_LONG_RELAXED(r->len) <= 0) {
                 // Jump over END_FOR instruction.
                 JUMPBY(oparg + 1);
                 DISPATCH();
@@ -3134,16 +3190,19 @@ dummy_func(
         op(_GUARD_NOT_EXHAUSTED_RANGE, (iter -- iter)) {
             _PyRangeIterObject *r = (_PyRangeIterObject *)PyStackRef_AsPyObjectBorrow(iter);
             assert(Py_TYPE(r) == &PyRangeIter_Type);
-            EXIT_IF(r->len <= 0);
+            EXIT_IF(FT_ATOMIC_LOAD_LONG_RELAXED(r->len) <= 0);
         }
 
         op(_ITER_NEXT_RANGE, (iter -- iter, next)) {
             _PyRangeIterObject *r = (_PyRangeIterObject *)PyStackRef_AsPyObjectBorrow(iter);
             assert(Py_TYPE(r) == &PyRangeIter_Type);
+#ifndef Py_GIL_DISABLED
             assert(r->len > 0);
-            long value = r->start;
-            r->start = value + r->step;
-            r->len--;
+#endif
+            long value = FT_ATOMIC_LOAD_LONG_RELAXED(r->start);
+            FT_ATOMIC_STORE_LONG_RELAXED(r->start, value + r->step);
+            FT_ATOMIC_STORE_LONG_RELAXED(r->len,
+                FT_ATOMIC_LOAD_LONG_RELAXED(r->len) - 1);
             PyObject *res = PyLong_FromLong(value);
             ERROR_IF(res == NULL, error);
             next = PyStackRef_FromPyObjectSteal(res);