Add specialization for range iterators and generators, both about as

thread-safe as without spcialization (i.e. not much to none at all).
python · Yhg1s · Mar 12, 2025 · Jan 13, 2025 · Jan 13, 2025 · Jan 13, 2025
commit 1433cd3e47f62c7bdda8abcba7b66a1345c14414
diff --git a/Objects/rangeobject.c b/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;
 }
 

@@ -3179,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();
@@ -3190,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);

diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
diff --git a/Python/specialize.c b/Python/specialize.c
@@ -2644,8 +2644,8 @@
 }
 #endif   // Py_STATS

 void
 _Py_Specialize_ForIter(_PyStackRef iter, _Py_CODEUNIT *instr, int oparg)
 {
    assert(ENABLE_SPECIALIZATION_FT);
    assert(_PyOpcode_Caches[FOR_ITER] == INLINE_CACHE_ENTRIES_FOR_ITER);
@@ -2668,12 +2668,13 @@
         specialize(instr, FOR_ITER_TUPLE);
         return;
     }
-#ifndef Py_GIL_DISABLED
     else if (tp == &PyRangeIter_Type) {
         specialize(instr, FOR_ITER_RANGE);
         return;
     }
     else if (tp == &PyGen_Type && oparg <= SHRT_MAX) {
+        // Generators are very much not thread-safe, so don't worry about
+        // the specialization not being thread-safe.
         assert(instr[oparg + INLINE_CACHE_ENTRIES_FOR_ITER + 1].op.code == END_FOR  ||
             instr[oparg + INLINE_CACHE_ENTRIES_FOR_ITER + 1].op.code == INSTRUMENTED_END_FOR
         );
@@ -2686,7 +2687,6 @@
         specialize(instr, FOR_ITER_GEN);
         return;
     }
-#endif
     SPECIALIZATION_FAIL(FOR_ITER,
                         _PySpecialization_ClassifyIterator(iter_o));
     unspecialize(instr);