optimized math

Author: tom-pytel

Include/internal/pycore_range.h

@@ -10,13 +10,25 @@ extern "C" {
 
 typedef struct {
     PyObject_HEAD
-    long start;
-    long step;
     long len;
+    long end;
+    long step;
     long stop;
 } _PyRangeIterObject;
 
-extern long _PyRangeIter_GetLength(_PyRangeIterObject *r, long start);
+static inline long
+_PyRangeIter_GetLengthAndStart(_PyRangeIterObject *r, long *value)
+{
+    long len = FT_ATOMIC_LOAD_LONG_RELAXED(r->len);
+    *value = r->end - r->step * len;
+    return len;
+}
+
+static inline void
+_PyRangeIter_SetLength(_PyRangeIterObject *r, long len)
+{
+    FT_ATOMIC_STORE_LONG_RELAXED(r->len, len);
+}
 
 #ifdef __cplusplus
 }

Objects/rangeobject.c

@@ -830,13 +830,10 @@ static PyObject *
 rangeiter_next(PyObject *op)
 {
     _PyRangeIterObject *r = (_PyRangeIterObject*)op;
-    long start = FT_ATOMIC_LOAD_LONG_RELAXED(r->start);
-    long len = _PyRangeIter_GetLength(r, start);
+    long start;
+    long len = _PyRangeIter_GetLengthAndStart(r, &start);
     if (len > 0) {
-#ifndef Py_GIL_DISABLED
-        r->len = len - 1;
-#endif
-        FT_ATOMIC_STORE_LONG_RELAXED(r->start, start + r->step);
+        _PyRangeIter_SetLength(r, len - 1);
         return PyLong_FromLong(start);
     }
     return NULL;
@@ -846,8 +843,9 @@ static PyObject *
 rangeiter_len(PyObject *op, PyObject *Py_UNUSED(ignored))
 {
     _PyRangeIterObject *r = (_PyRangeIterObject*)op;
-    long start = FT_ATOMIC_LOAD_LONG_RELAXED(r->start);
-    return PyLong_FromLong(_PyRangeIter_GetLength(r, start));
+    long start;
+    long len = _PyRangeIter_GetLengthAndStart(r, &start);
+    return PyLong_FromLong(len);
 }
 
 PyDoc_STRVAR(length_hint_doc,
@@ -861,7 +859,9 @@ rangeiter_reduce(PyObject *op, PyObject *Py_UNUSED(ignored))
     PyObject *range;
 
     /* create a range object for pickling */
-    start = PyLong_FromLong(FT_ATOMIC_LOAD_LONG_RELAXED(r->start));
+    long lstart;
+    (void)_PyRangeIter_GetLengthAndStart(r, &lstart);
+    start = PyLong_FromLong(lstart);
     if (start == NULL)
         goto err;
     stop = PyLong_FromLong(r->stop);
@@ -892,14 +892,13 @@ rangeiter_setstate(PyObject *op, PyObject *state)
     if (index == -1 && PyErr_Occurred())
         return NULL;
     /* silently clip the index value */
-    long start = FT_ATOMIC_LOAD_LONG_RELAXED(r->start);
-    long len = _PyRangeIter_GetLength(r, start);
+    long start;
+    long len = _PyRangeIter_GetLengthAndStart(r, &start);
     if (index < 0)
         index = 0;
     else if (index > len)
         index = len; /* exhausted iterator */
-    r->len = len - index;
-    FT_ATOMIC_STORE_LONG_RELAXED(r->start, start + index * r->step);
+    _PyRangeIter_SetLength(r, len - index);
     Py_RETURN_NONE;
 }
 
@@ -983,9 +982,9 @@ fast_range_iter(long start, long stop, long step, long len)
     _PyRangeIterObject *it = PyObject_New(_PyRangeIterObject, &PyRangeIter_Type);
     if (it == NULL)
         return NULL;
-    it->start = start;
-    it->step = step;
     it->len = len;
+    it->end = start + step * len;
+    it->step = step;
     it->stop = stop;
     return (PyObject *)it;
 }

Python/bytecodes.c

@@ -3284,25 +3284,14 @@ dummy_func(
             _PyRangeIterObject *r = (_PyRangeIterObject *)PyStackRef_AsPyObjectBorrow(iter);
             assert(Py_TYPE(r) == &PyRangeIter_Type);
             STAT_INC(FOR_ITER, hit);
-#ifdef Py_GIL_DISABLED
-            long value = FT_ATOMIC_LOAD_LONG_RELAXED(r->start);
-            long len = _PyRangeIter_GetLength(r, value);
+            long value;
+            long len = _PyRangeIter_GetLengthAndStart(r, &value);
             if (len <= 0) {
                 // Jump over END_FOR instruction.
                 JUMPBY(oparg + 1);
                 DISPATCH();
             }
-            FT_ATOMIC_STORE_LONG_RELAXED(r->start, value + r->step);
-#else  // the code above will work for GIL build but below is faster
-            if (r->len <= 0) {
-                // Jump over END_FOR instruction.
-                JUMPBY(oparg + 1);
-                DISPATCH();
-            }
-            long value = r->start;
-            r->start = value + r->step;
-            r->len--;
-#endif
+            _PyRangeIter_SetLength(r, len - 1);
             PyObject *res = PyLong_FromLong(value);
             ERROR_IF(res == NULL, error);
             next = PyStackRef_FromPyObjectSteal(res);
@@ -3312,17 +3301,10 @@ dummy_func(
         op(_ITER_NEXT_RANGE_TIER_TWO, (iter -- iter, next)) {
             _PyRangeIterObject *r = (_PyRangeIterObject *)PyStackRef_AsPyObjectBorrow(iter);
             assert(Py_TYPE(r) == &PyRangeIter_Type);
-#ifdef Py_GIL_DISABLED
-            long value = FT_ATOMIC_LOAD_LONG_RELAXED(r->start);
-            long len = _PyRangeIter_GetLength(r, value);
+            long value;
+            long len = _PyRangeIter_GetLengthAndStart(r, &value);
             EXIT_IF(len <= 0);
-            FT_ATOMIC_STORE_LONG_RELAXED(r->start, value + r->step);
-#else  // the code above will work for GIL build but below is faster
-            EXIT_IF(r->len <= 0);
-            long value = r->start;
-            r->start = value + r->step;
-            r->len--;
-#endif
+            _PyRangeIter_SetLength(r, len - 1);
             PyObject *res = PyLong_FromLong(value);
             ERROR_IF(res == NULL, error);
             next = PyStackRef_FromPyObjectSteal(res);

Python/executor_cases.c.h

@@ -674,7 +674,8 @@ def has_error_without_pop(op: parser.CodeDef) -> bool:
     "JUMP_TO_LABEL",
     "restart_backoff_counter",
     "_Py_ReachedRecursionLimit",
-    "_PyRangeIter_GetLength",
+    "_PyRangeIter_GetLengthAndStart",
+    "_PyRangeIter_SetLength",
 )
 
 def find_stmt_start(node: parser.CodeDef, idx: int) -> lexer.Token: