GH-132554: Specialize GET_ITER and FOR_ITER for range

Include/cpython/object.h

@@ -239,6 +239,9 @@ struct _typeobject {
      * Otherwise, limited to MAX_VERSIONS_PER_CLASS (defined elsewhere).
      */
     uint16_t tp_versions_used;
+    /* Returns the object at the index given, or NULL if out-of-bounds
+     * Never raises an exception. */
+    iterindexfunc tp_iterindex;
 };
 
 #define _Py_ATTR_CACHE_UNUSED (30000)  // (see tp_versions_used)

Include/internal/pycore_code.h

@@ -313,6 +313,7 @@ extern void _Py_Specialize_CompareOp(_PyStackRef lhs, _PyStackRef rhs,
                                      _Py_CODEUNIT *instr, int oparg);
 extern void _Py_Specialize_UnpackSequence(_PyStackRef seq, _Py_CODEUNIT *instr,
                                           int oparg);
+extern void _Py_Specialize_GetIter(_PyStackRef iter, _Py_CODEUNIT *instr);
 extern void _Py_Specialize_ForIter(_PyStackRef iter, _PyStackRef null_or_index, _Py_CODEUNIT *instr, int oparg);
 extern void _Py_Specialize_Send(_PyStackRef receiver, _Py_CODEUNIT *instr);
 extern void _Py_Specialize_ToBool(_PyStackRef value, _Py_CODEUNIT *instr);

Include/internal/pycore_magic_number.h

@@ -281,6 +281,7 @@ Known values:
     Python 3.15a1 3651 (Simplify LOAD_CONST)
     Python 3.15a1 3652 (Virtual iterators)
     Python 3.15a1 3653 (Fix handling of opcodes that may leave operands on the stack when optimizing LOAD_FAST)
+    Python 3.15a1 3654 (Specialize GET_ITER. Add FOR_ITER_INDEX)
 
 
     Python 3.16 will start with 3700
@@ -294,7 +295,7 @@ PC/launcher.c must also be updated.
 
 */
 
-#define PYC_MAGIC_NUMBER 3653
+#define PYC_MAGIC_NUMBER 3654
 /* This is equivalent to converting PYC_MAGIC_NUMBER to 2 bytes
    (little-endian) and then appending b'\r\n'. */
 #define PYC_MAGIC_NUMBER_TOKEN \

Include/internal/pycore_range.h

@@ -8,13 +8,52 @@ extern "C" {
 #  error "this header requires Py_BUILD_CORE define"
 #endif
 
+typedef struct {
+    PyObject_HEAD
+    PyObject *start;
+    PyObject *stop;
+    PyObject *step;
+    PyObject *length;
+} rangeobject;
+
 typedef struct {
     PyObject_HEAD
     long start;
     long step;
     long len;
 } _PyRangeIterObject;
 
+// Does this range have step == 1 and both start and stop in compact int range?
+static inline int
+_PyRange_IsSimpleCompact(PyObject *range) {
+    assert(PyRange_Check(range));
+    rangeobject *r = (rangeobject*)range;
+    if (_PyLong_IsCompact((PyLongObject *)r->start) &&
+        _PyLong_IsCompact((PyLongObject *)r->stop) &&
+        r->step == _PyLong_GetOne()
+    ) {
+        return 1;
+    }
+    return 0;
+}
+
+static inline 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);
+}
+
+static inline Py_ssize_t
+_PyRange_GetStopIfCompact(PyObject *range) {
+    assert(PyRange_Check(range));
+    rangeobject *r = (rangeobject*)range;
+    assert(_PyLong_IsCompact((PyLongObject *)r->stop));
+    return _PyLong_CompactValue((PyLongObject *)r->stop);
+}
+
 #ifdef __cplusplus
 }
 #endif

Include/internal/pycore_stackref.h

@@ -115,6 +115,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)
 {
@@ -329,10 +337,18 @@ static inline _PyStackRef
 PyStackRef_IncrementTaggedIntNoOverflow(_PyStackRef ref)
 {
     assert((ref.bits & Py_TAG_BITS) == Py_INT_TAG); // Is tagged int
-    assert((ref.bits & (~Py_TAG_BITS)) != (INT_MAX & (~Py_TAG_BITS))); // Isn't about to overflow
+    assert((ref.bits & (~Py_TAG_BITS)) != (INTPTR_MAX & (~Py_TAG_BITS))); // Isn't about to overflow
     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)
 
 #ifdef Py_GIL_DISABLED
@@ -849,6 +865,20 @@ _Py_TryXGetStackRef(PyObject **src, _PyStackRef *out)
         }                                                               \
     } while (0)
 
+
+static inline _PyStackRef
+PyStackRef_BoxInt(_PyStackRef i)
+{
+    assert(PyStackRef_IsTaggedInt(i));
+    intptr_t val = PyStackRef_UntagInt(i);
+    PyObject *boxed = PyLong_FromSsize_t(val);
+    if (boxed == NULL) {
+        return PyStackRef_ERROR;
+    }
+    return PyStackRef_FromPyObjectSteal(boxed);
+}
+
+
 #ifdef __cplusplus
 }
 #endif