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

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

Include/internal/pycore_magic_number.h

@@ -286,6 +286,7 @@ Known values:
  Python 3.15a1 3653 (Fix handling of opcodes that may leave operands on the stack when optimizing LOAD_FAST)
  Python 3.15a1 3654 (Fix missing exception handlers in logical expression)
  Python 3.15a1 3655 (Fix miscompilation of some module-level annotations)
+ Python 3.15a1 3656 (Specialize GET_ITER. Add FOR_ITER_INDEX)
 
 
  Python 3.16 will start with 3700
@@ -299,7 +300,7 @@ PC/launcher.c must also be updated.
 
 */
 
-#define PYC_MAGIC_NUMBER 3655
+#define PYC_MAGIC_NUMBER 3656
 /* 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
@@ -862,6 +878,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