Re-use of a PyLong for integer range iterator

- By not allocating a new PyLong on each iteration, we save significant
  time; in a benchmark with a large range a bit more than a factor of 2,
  with a small range a few percent (~8% with only interned ints)
- To make sure this doesn't break when assigning to different variables
  using next(), we have to decrement the target object's reference counter
  before calling iternext(). We do this by checking if the next
  instruction is a STORE_FAST, and treat the FOR_ITER/STORE_FAST pair as
  a sort of super-instruction, saving some pushing and popping from the
  stack as well.
- We do have to be careful with DECREFing before setting the new local,
  as mentioned in the comment to the SETLOCAL macro:
  'This is because it is possible that during the DECREF the frame is
   accessed by other code (e.g. a __del__ method or gc.collect()) and the
   variable would be pointing to already-freed memory.'
  Therefore we restrict this only to long types with a ref count >1
- One caveat is that range will no longer result in interned integers,
  so 'for i in range(10)' will result in 'i is 5' always False.

Author: heeres

Objects/rangeobject.c

@@ -770,16 +770,45 @@ typedef struct {
         long    start;
         long    step;
         long    len;
+        PyObject *fast_obj;
 } rangeiterobject;
 
 static PyObject *
 rangeiter_next(rangeiterobject *r)
 {
-    if (r->index < r->len)
+    if (r->index < r->len) {
         /* cast to unsigned to avoid possible signed overflow
            in intermediate calculations. */
-        return PyLong_FromLong((long)(r->start +
-                                      (unsigned long)(r->index++) * r->step));
+        long cur_val = (long)(r->start + (unsigned long)(r->index++) * r->step);
+
+        // If we are the only one holding a reference, modify object in place
+        if (Py_REFCNT(r->fast_obj) == 1) {
+            long abs_val, abs_val_msb;
+            int sign;
+            PyLongObject *l = (PyLongObject*)r->fast_obj;
+            Py_INCREF(r->fast_obj);
+            if (cur_val < 0) {
+                abs_val = 0U - (unsigned long)cur_val;
+                sign = -1;
+            }
+            else {
+                abs_val = cur_val;
+                sign = cur_val ? 1 : 0;
+            }
+            abs_val_msb = abs_val >> PYLONG_BITS_IN_DIGIT;
+            if (!abs_val_msb)
+                l->ob_digit[0] = Py_SAFE_DOWNCAST(abs_val, unsigned long, digit);
+            else {
+                sign *= 2;
+                l->ob_digit[0] = Py_SAFE_DOWNCAST(abs_val&PyLong_MASK, unsigned long, digit);
+                l->ob_digit[1] = Py_SAFE_DOWNCAST(abs_val_msb, unsigned long, digit);
+            }
+            Py_SET_SIZE(l, sign);
+            return r->fast_obj;
+        }
+        else
+            return PyLong_FromLong(cur_val);
+    }
     return NULL;
 }
 
@@ -837,6 +866,13 @@ rangeiter_setstate(rangeiterobject *r, PyObject *state)
     Py_RETURN_NONE;
 }
 
+static void
+rangeiter_dealloc(rangeiterobject *r)
+{
+    Py_XDECREF(r->fast_obj);
+    PyObject_Free(r);
+}
+
 PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");
 PyDoc_STRVAR(setstate_doc, "Set state information for unpickling.");
 
@@ -856,7 +892,7 @@ PyTypeObject PyRangeIter_Type = {
         sizeof(rangeiterobject),                /* tp_basicsize */
         0,                                      /* tp_itemsize */
         /* methods */
-        (destructor)PyObject_Del,               /* tp_dealloc */
+        (destructor)rangeiter_dealloc,          /* tp_dealloc */
         0,                                      /* tp_vectorcall_offset */
         0,                                      /* tp_getattr */
         0,                                      /* tp_setattr */
@@ -921,6 +957,16 @@ fast_range_iter(long start, long stop, long step)
     unsigned long ulen;
     if (it == NULL)
         return NULL;
+
+    // Allocate fast return object; make sure it's a 2 digit PyLong
+    // We hold a reference that we have to release in _dealloc
+    it->fast_obj = (PyObject*)PyLong_FromLong(1L<<(PYLONG_BITS_IN_DIGIT+1));
+    if (it->fast_obj == NULL)
+    {
+        Py_DECREF(it);
+        return NULL;
+    }
+
     it->start = start;
     it->step = step;
     ulen = get_len_of_range(start, stop, step);

Python/ceval.c

@@ -1842,7 +1842,6 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, PyFrameObject *f, int throwflag)
         }
 
         case TARGET(STORE_FAST): {
-            PREDICTED(STORE_FAST);
             PyObject *value = POP();
             SETLOCAL(oparg, value);
             DISPATCH();
@@ -3995,14 +3994,38 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, PyFrameObject *f, int throwflag)
         case TARGET(FOR_ITER): {
             PREDICTED(FOR_ITER);
             /* before: [iter]; after: [iter, iter()] *or* [] */
-            PyObject *iter = TOP();
-            PyObject *next = (*Py_TYPE(iter)->tp_iternext)(iter);
+            PyObject *cur_local, *next, *iter = TOP();
+
+            // If FOR_ITER is followed by STORE_FAST, pre-release a reference on PyLong before
+            // calling iternext, indicating that the target object can be reused (e.g. for range())
+            int for_target, for_store_fast = _Py_OPCODE(*next_instr) == STORE_FAST ? 1 : 0;
+            if (for_store_fast) {
+                for_target = _Py_OPARG(*next_instr);
+                cur_local = GETLOCAL(for_target);
+                // Can't do this when ref count == 1, because an object should persist after last iteration
+                if (cur_local && PyLong_Check(cur_local) && Py_REFCNT(cur_local)>1) {
+                    Py_DECREF(cur_local);
+                    for_store_fast = 2;         // Indicate reference count was already decremented
+                }
+            }
+            next = (*Py_TYPE(iter)->tp_iternext)(iter);
             if (next != NULL) {
-                PUSH(next);
-                PREDICT(STORE_FAST);
-                PREDICT(UNPACK_SEQUENCE);
+                if (for_store_fast) {
+                    if (for_store_fast == 2)    // PyLong with ref count that was >1, no need for additional book keeping
+                        GETLOCAL(for_target) = next;
+                    else
+                        SETLOCAL(for_target, next);
+                    next_instr++;               // STORE_FAST already performed
+                }
+                else {
+                    PUSH(next);
+                    PREDICT(UNPACK_SEQUENCE);
+                }
                 DISPATCH();
             }
+            // Restore reference count of local if it was pre-decremented
+            if (for_store_fast == 2)
+                Py_INCREF(cur_local);
             if (_PyErr_Occurred(tstate)) {
                 if (!_PyErr_ExceptionMatches(tstate, PyExc_StopIteration)) {
                     goto error;