1
1
#include "Python.h"
2
- #include "pycore_call.h" // _PyObject_CallNoArgs()
3
- #include "pycore_ceval.h" // _PyEval_GetBuiltin()
4
- #include "pycore_long.h" // _PyLong_GetZero()
5
- #include "pycore_moduleobject.h" // _PyModule_GetState()
6
- #include "pycore_typeobject.h" // _PyType_GetModuleState()
7
- #include "pycore_object.h" // _PyObject_GC_TRACK()
8
- #include "pycore_tuple.h" // _PyTuple_ITEMS()
2
+ #include "pycore_call.h" // _PyObject_CallNoArgs()
3
+ #include "pycore_ceval.h" // _PyEval_GetBuiltin()
4
+ #include "pycore_critical_section.h" // Py_BEGIN_CRITICAL_SECTION()
5
+ #include "pycore_long.h" // _PyLong_GetZero()
6
+ #include "pycore_moduleobject.h" // _PyModule_GetState()
7
+ #include "pycore_typeobject.h" // _PyType_GetModuleState()
8
+ #include "pycore_object.h" // _PyObject_GC_TRACK()
9
+ #include "pycore_tuple.h" // _PyTuple_ITEMS()
9
10
10
- #include <stddef.h> // offsetof()
11
+ #include <stddef.h> // offsetof()
11
12
12
13
/* Itertools module written and maintained
13
14
by Raymond D. Hettinger <python@rcn.com>
@@ -3254,7 +3255,7 @@ fast_mode: when cnt an integer < PY_SSIZE_T_MAX and no step is specified.
3254
3255
3255
3256
assert(cnt != PY_SSIZE_T_MAX && long_cnt == NULL && long_step==PyLong(1));
3256
3257
Advances with: cnt += 1
3257
- When count hits Y_SSIZE_T_MAX , switch to slow_mode.
3258
+ When count hits PY_SSIZE_T_MAX , switch to slow_mode.
3258
3259
3259
3260
slow_mode: when cnt == PY_SSIZE_T_MAX, step is not int(1), or cnt is a float.
3260
3261
@@ -3386,7 +3387,7 @@ count_nextlong(countobject *lz)
3386
3387
3387
3388
long_cnt = lz -> long_cnt ;
3388
3389
if (long_cnt == NULL ) {
3389
- /* Switch to slow_mode */
3390
+ /* Switching from fast mode */
3390
3391
long_cnt = PyLong_FromSsize_t (PY_SSIZE_T_MAX );
3391
3392
if (long_cnt == NULL )
3392
3393
return NULL ;
@@ -3403,9 +3404,35 @@ count_nextlong(countobject *lz)
3403
3404
static PyObject *
3404
3405
count_next (countobject * lz )
3405
3406
{
3406
- if (lz -> cnt == PY_SSIZE_T_MAX )
3407
- return count_nextlong (lz );
3408
- return PyLong_FromSsize_t (lz -> cnt ++ );
3407
+ PyObject * returned ;
3408
+ Py_ssize_t cnt ;
3409
+
3410
+ cnt = FT_ATOMIC_LOAD_SSIZE_RELAXED (lz -> cnt );
3411
+ for (;;) {
3412
+ if (cnt == PY_SSIZE_T_MAX ) {
3413
+ /* slow mode */
3414
+ Py_BEGIN_CRITICAL_SECTION (lz );
3415
+ returned = count_nextlong (lz );
3416
+ Py_END_CRITICAL_SECTION ();
3417
+ return returned ;
3418
+ }
3419
+ #ifdef Py_GIL_DISABLED
3420
+ /* thread-safe fast version (increment by one).
3421
+ * If lz->cnt changed between the pervious read and now,
3422
+ * that means another thread got in our way. In this case,
3423
+ * update cnt to new value of lz->cnt, and try again.
3424
+ * Otherwise, (no other thread updated lz->cnt),
3425
+ * atomically update lz->cnt with the incremented value and
3426
+ * then return cnt (the previous value)
3427
+ */
3428
+ if (_Py_atomic_compare_exchange_ssize (& lz -> cnt , & cnt , cnt + 1 )) {
3429
+ return PyLong_FromSsize_t (cnt );
3430
+ }
3431
+ #else
3432
+ /* fast mode when GIL is enabled */
3433
+ return PyLong_FromSsize_t (lz -> cnt ++ );
3434
+ #endif
3435
+ }
3409
3436
}
3410
3437
3411
3438
static PyObject *
0 commit comments