Allow range starts other than zero

Author: markshannon

Include/internal/pycore_range.h

@@ -17,6 +17,7 @@ typedef struct {
 
 // Does this range have start == 0, step == 1 and step in compact int range?
 int _PyRange_IsSimpleCompact(PyObject *range);
+Py_ssize_t _PyRange_GetStartIfCompact(PyObject *range);
 Py_ssize_t _PyRange_GetStopIfCompact(PyObject *range);
 
 #ifdef __cplusplus

Include/internal/pycore_stackref.h

@@ -101,6 +101,14 @@ PyStackRef_IsTaggedInt(_PyStackRef ref)
     return (ref.index & 1) == 1;
 }
 
+static inline bool
+PyStackRef_TaggedIntLessThan(_PyStackRef a, _PyStackRef b)
+{
+    assert(PyStackRef_IsTaggedInt(a));
+    assert(PyStackRef_IsTaggedInt(b));
+    return ((intptr_t)a.bits) < ((intptr_t)b.bits);
+}
+
 static inline PyObject *
 _PyStackRef_AsPyObjectBorrow(_PyStackRef ref, const char *filename, int linenumber)
 {
@@ -274,6 +282,14 @@ PyStackRef_IncrementTaggedIntNoOverflow(_PyStackRef ref)
     return (_PyStackRef){ .bits = ref.bits + 4 };
 }
 
+static inline bool
+PyStackRef_TaggedIntLessThan(_PyStackRef a, _PyStackRef b)
+{
+    assert(PyStackRef_IsTaggedInt(a));
+    assert(PyStackRef_IsTaggedInt(b));
+    return ((intptr_t)a.bits) < ((intptr_t)b.bits);
+}
+
 #define PyStackRef_IsDeferredOrTaggedInt(ref) (((ref).bits & Py_TAG_REFCNT) != 0)
 
 extern _PyStackRef PyStackRef_BoxInt(_PyStackRef i);

Lib/test/test_opcache.py

@@ -1832,15 +1832,6 @@ def for_iter_tuple():
         self.assert_specialized(for_iter_tuple, "FOR_ITER_TUPLE")
         self.assert_no_opcode(for_iter_tuple, "FOR_ITER")
 
-        r = range(10)
-        def for_iter_range():
-            for i in r:
-                self.assertIn(i, r)
-
-        for_iter_range()
-        self.assert_specialized(for_iter_range, "FOR_ITER_RANGE")
-        self.assert_no_opcode(for_iter_range, "FOR_ITER")
-
         def for_iter_generator():
             for i in (i for i in range(10)):
                 i + 1
@@ -1849,6 +1840,44 @@ def for_iter_generator():
         self.assert_specialized(for_iter_generator, "FOR_ITER_GEN")
         self.assert_no_opcode(for_iter_generator, "FOR_ITER")
 
+    @cpython_only
+    @requires_specialization_ft
+    def test_for_iter_range(self):
+        r = range(10)
+        def for_iter1():
+            for i in r:
+                self.assertIn(i, r)
+
+        for_iter1()
+        self.assert_specialized(for_iter1, "FOR_ITER_RANGE")
+        self.assert_no_opcode(for_iter1, "FOR_ITER")
+
+        r = range(-10, 0)
+        def for_iter2():
+            for i in r:
+                self.assertIn(i, r)
+
+        for_iter2()
+        self.assert_specialized(for_iter2, "FOR_ITER_RANGE")
+        self.assert_no_opcode(for_iter2, "FOR_ITER")
+
+        r = range(100_000, 100_010)
+        def for_iter3():
+            for i in r:
+                self.assertIn(i, r)
+
+        for_iter3()
+        self.assert_specialized(for_iter3, "FOR_ITER_RANGE")
+        self.assert_no_opcode(for_iter3, "FOR_ITER")
+
+        r = range((1 << 65), (1 << 65)+10)
+        def for_iter4():
+            for i in r:
+                self.assertIn(i, r)
+
+        for_iter4()
+        self.assert_specialized(for_iter4, "FOR_ITER")
+
 
 if __name__ == "__main__":
     unittest.main()

Objects/rangeobject.c

@@ -160,20 +160,29 @@ int
 _PyRange_IsSimpleCompact(PyObject *range) {
     assert(PyRange_Check(range));
     rangeobject *r = (rangeobject*)range;
-    if (r->start == _PyLong_GetZero() && r->step == _PyLong_GetOne() &&
-        _PyLong_IsNonNegativeCompact((PyLongObject *)r->stop)
+    if (_PyLong_IsCompact((PyLongObject *)r->start) &&
+        _PyLong_IsCompact((PyLongObject *)r->stop) &&
+        r->step == _PyLong_GetOne()
     ) {
         return 1;
     }
     return 0;
 }
 
+Py_ssize_t
+_PyRange_GetStartIfCompact(PyObject *range) {
+    assert(PyRange_Check(range));
+    rangeobject *r = (rangeobject*)range;
+    assert(_PyLong_IsCompact((PyLongObject *)r->start));
+    return _PyLong_CompactValue((PyLongObject *)r->start);
+}
+
 Py_ssize_t
 _PyRange_GetStopIfCompact(PyObject *range) {
     assert(PyRange_Check(range));
     rangeobject *r = (rangeobject*)range;
-    assert(_PyLong_IsNonNegativeCompact((PyLongObject *)r->stop));
-    return _PyLong_GetNonNegativeCompactValue((PyLongObject *)r->stop);
+    assert(_PyLong_IsCompact((PyLongObject *)r->stop));
+    return _PyLong_CompactValue((PyLongObject *)r->stop);
 }
 
 PyDoc_STRVAR(range_doc,

Python/bytecodes.c

@@ -3080,10 +3080,11 @@ dummy_func(
             else {
                 PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iterable);
                 if (tp == &PyRange_Type && _PyRange_IsSimpleCompact(iter_o)) {
+                    Py_ssize_t start = _PyRange_GetStartIfCompact(iter_o);
                     Py_ssize_t stop = _PyRange_GetStopIfCompact(iter_o);
                     PyStackRef_CLOSE(iterable);
                     iter = PyStackRef_TagInt(stop);
-                    index_or_null = PyStackRef_TagInt(0);
+                    index_or_null = PyStackRef_TagInt(start);
                 }
                 else {
                     iter_o =  PyObject_GetIter(iter_o);
@@ -3171,8 +3172,7 @@ dummy_func(
             /* before: [iter]; after: [iter, iter()] *or* [] (and jump over END_FOR.) */
             if (PyStackRef_IsTaggedInt(null_or_index)) {
                 if (PyStackRef_IsTaggedInt(iter)) {
-                    if (PyStackRef_Is(iter, null_or_index)) {
-                        null_or_index = PyStackRef_TagInt(-1);
+                    if (!PyStackRef_TaggedIntLessThan(null_or_index, iter)) {
                         JUMPBY(oparg + 1);
                         DISPATCH();
 
@@ -3242,14 +3242,11 @@ dummy_func(
 
 
         inst(INSTRUMENTED_FOR_ITER, (unused/1, iter, null_or_index -- iter, null_or_index, next)) {
-            PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
             if (PyStackRef_IsTaggedInt(null_or_index)) {
                 if (PyStackRef_IsTaggedInt(iter)) {
-                    if (PyStackRef_Is(iter, null_or_index)) {
-                        null_or_index = PyStackRef_TagInt(-1);
+                    if (!PyStackRef_TaggedIntLessThan(null_or_index, iter)) {
                         JUMPBY(oparg + 1);
                         DISPATCH();
-
                     }
                     next = PyStackRef_BoxInt(null_or_index);
                     if (PyStackRef_IsNull(next)) {
@@ -3258,6 +3255,7 @@ dummy_func(
                     null_or_index = PyStackRef_IncrementTaggedIntNoOverflow(null_or_index);
                 }
                 else {
+                    PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
                     next = _PyForIter_NextWithIndex(iter_o, null_or_index);
                     if (PyStackRef_IsNull(next)) {
                         JUMPBY(oparg + 1);
@@ -3268,6 +3266,7 @@ dummy_func(
                 INSTRUMENTED_JUMP(this_instr, next_instr, PY_MONITORING_EVENT_BRANCH_LEFT);
             }
             else {
+                PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
                 PyObject *next_o = (*Py_TYPE(iter_o)->tp_iternext)(iter_o);
                 if (next_o != NULL) {
                     next = PyStackRef_FromPyObjectSteal(next_o);
@@ -3423,7 +3422,7 @@ dummy_func(
         }
 
         replaced op(_ITER_JUMP_RANGE, (iter, null_or_index -- iter, null_or_index)) {
-            if (PyStackRef_Is(iter, null_or_index)) {
+            if (!PyStackRef_TaggedIntLessThan(null_or_index, iter)) {
                 // Jump over END_FOR instruction.
                 JUMPBY(oparg + 1);
                 DISPATCH();
@@ -3432,7 +3431,7 @@ dummy_func(
 
         // Only used by Tier 2
         op(_GUARD_NOT_EXHAUSTED_RANGE, (iter, null_or_index -- iter, null_or_index)) {
-            EXIT_IF(PyStackRef_Is(iter, null_or_index));
+            EXIT_IF(!PyStackRef_TaggedIntLessThan(null_or_index, iter));
         }
 
         op(_ITER_NEXT_RANGE, (iter, null_or_index -- iter, null_or_index, next)) {

Python/executor_cases.c.h

@@ -2912,6 +2912,7 @@ _Py_Specialize_ForIter(_PyStackRef iter, _PyStackRef null_or_index, _Py_CODEUNIT
 
     if (PyStackRef_IsNull(null_or_index)) {
 #ifdef Py_GIL_DISABLED
+        PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
         // Only specialize for uniquely referenced iterators, so that we know
         // they're only referenced by this one thread. This is more limiting
         // than we need (even `it = iter(mylist); for item in it:` won't get

Python/generated_cases.c.h

@@ -202,6 +202,14 @@ _PyStackRef PyStackRef_TagInt(intptr_t i)
     return (_PyStackRef){ .index = (i << 1) + 1 };
 }
 
+bool
+PyStackRef_TaggedIntLessThan(_PyStackRef a, _PyStackRef b)
+{
+    assert(PyStackRef_IsTaggedInt(a));
+    assert(PyStackRef_IsTaggedInt(b));
+    return a.bits < b.bits;
+}
+
 intptr_t
 PyStackRef_UntagInt(_PyStackRef i)
 {

Python/specialize.c

@@ -685,6 +685,7 @@ def has_error_without_pop(op: parser.CodeDef) -> bool:
     "_PyRange_GetStopIfCompact",
     "PyStackRef_BoxInt",
     "PyStackRef_TYPE",
+    "PyStackRef_TaggedIntLessThan",
 )