GH-132554: "Virtual" iterators

Include/internal/pycore_ceval.h

@@ -353,6 +353,8 @@ PyAPI_FUNC(_PyStackRef) _PyFloat_FromDouble_ConsumeInputs(_PyStackRef left, _PyS
 extern int _PyRunRemoteDebugger(PyThreadState *tstate);
 #endif
 
+_PyStackRef _PyForIter_NextWithIndex(PyObject *seq, _PyStackRef index);
+
 #ifdef __cplusplus
 }
 #endif

Include/internal/pycore_code.h

@@ -313,7 +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_ForIter(_PyStackRef iter, _Py_CODEUNIT *instr, int oparg);
+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);
 extern void _Py_Specialize_ContainsOp(_PyStackRef value, _Py_CODEUNIT *instr);

Include/internal/pycore_compile.h

@@ -95,6 +95,7 @@ typedef enum {
 enum _PyCompile_FBlockType {
  COMPILE_FBLOCK_WHILE_LOOP,
  COMPILE_FBLOCK_FOR_LOOP,
+ COMPILE_FBLOCK_ASYNC_FOR_LOOP,
  COMPILE_FBLOCK_TRY_EXCEPT,
  COMPILE_FBLOCK_FINALLY_TRY,
  COMPILE_FBLOCK_FINALLY_END,

Include/internal/pycore_magic_number.h

@@ -279,6 +279,8 @@ Known values:
  Python 3.14b1 3624 (Don't optimize LOAD_FAST when local is killed by DELETE_FAST)
  Python 3.15a0 3650 (Initial version)
  Python 3.15a1 3651 (Simplify LOAD_CONST)
+ Python 3.15a1 3652 (Virtual iterators)
+
 
  Python 3.16 will start with 3700
 
@@ -291,7 +293,7 @@ PC/launcher.c must also be updated.
 
 */
 
-#define PYC_MAGIC_NUMBER 3651
+#define PYC_MAGIC_NUMBER 3652
 /* 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_stackref.h

@@ -232,18 +232,22 @@ extern intptr_t PyStackRef_UntagInt(_PyStackRef ref);
 
 extern _PyStackRef PyStackRef_TagInt(intptr_t i);
 
+/* Increments a tagged int, but does not check for overflow */
+extern _PyStackRef PyStackRef_IncrementTaggedIntNoOverflow(_PyStackRef ref);
+
 extern bool
 PyStackRef_IsNullOrInt(_PyStackRef ref);
 
 #else
 
 #define Py_INT_TAG 3
 #define Py_TAG_REFCNT 1
+#define Py_TAG_BITS 3
 
 static inline bool
 PyStackRef_IsTaggedInt(_PyStackRef i)
 {
- return (i.bits & Py_INT_TAG) == Py_INT_TAG;
+ return (i.bits & Py_TAG_BITS) == Py_INT_TAG;
 }
 
 static inline _PyStackRef
@@ -262,12 +266,21 @@ PyStackRef_UntagInt(_PyStackRef i)
 }
 
 
+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
+ return (_PyStackRef){ .bits = ref.bits + 4 };
+}
+
+#define PyStackRef_IsDeferredOrTaggedInt(ref) (((ref).bits & Py_TAG_REFCNT) != 0)
+
 #ifdef Py_GIL_DISABLED
 
 #define Py_TAG_DEFERRED Py_TAG_REFCNT
 
 #define Py_TAG_PTR ((uintptr_t)0)
-#define Py_TAG_BITS ((uintptr_t)1)
 
 
 static const _PyStackRef PyStackRef_NULL = { .bits = Py_TAG_DEFERRED};
@@ -379,7 +392,7 @@ PyStackRef_FromPyObjectBorrow(PyObject *obj)
  do { \
  _PyStackRef _close_tmp = (REF); \
  assert(!PyStackRef_IsNull(_close_tmp)); \
- if (!PyStackRef_IsDeferred(_close_tmp)) { \
+ if (!PyStackRef_IsDeferredOrTaggedInt(_close_tmp)) { \
  Py_DECREF(PyStackRef_AsPyObjectBorrow(_close_tmp)); \
  } \
  } while (0)
@@ -395,7 +408,7 @@ static inline _PyStackRef
 PyStackRef_DUP(_PyStackRef stackref)
 {
  assert(!PyStackRef_IsNull(stackref));
- if (PyStackRef_IsDeferred(stackref)) {
+ if (PyStackRef_IsDeferredOrTaggedInt(stackref)) {
  return stackref;
  }
  Py_INCREF(PyStackRef_AsPyObjectBorrow(stackref));
@@ -442,7 +455,6 @@ PyStackRef_AsStrongReference(_PyStackRef stackref)
 /* References to immortal objects always have their tag bit set to Py_TAG_REFCNT
  * as they can (must) have their reclamation deferred */
 
-#define Py_TAG_BITS 3
 #if _Py_IMMORTAL_FLAGS != Py_TAG_REFCNT
 # error "_Py_IMMORTAL_FLAGS != Py_TAG_REFCNT"
 #endif
@@ -678,50 +690,44 @@ PyStackRef_XCLOSE(_PyStackRef ref)
 
 #endif // !defined(Py_GIL_DISABLED) && defined(Py_STACKREF_DEBUG)
 
-#define PyStackRef_TYPE(stackref) Py_TYPE(PyStackRef_AsPyObjectBorrow(stackref))
+static inline PyTypeObject *
+PyStackRef_TYPE(_PyStackRef stackref) {
+ if (PyStackRef_IsTaggedInt(stackref)) {
+ return &PyLong_Type;
+ }
+ return Py_TYPE(PyStackRef_AsPyObjectBorrow(stackref));
+}
 
 // Converts a PyStackRef back to a PyObject *, converting the
 // stackref to a new reference.
 #define PyStackRef_AsPyObjectNew(stackref) Py_NewRef(PyStackRef_AsPyObjectBorrow(stackref))
 
 // StackRef type checks
 
-static inline bool
-PyStackRef_GenCheck(_PyStackRef stackref)
-{
- return PyGen_Check(PyStackRef_AsPyObjectBorrow(stackref));
-}
+#define STACKREF_CHECK_FUNC(T) \
+ static inline bool \
+ PyStackRef_ ## T ## Check(_PyStackRef stackref) { \
+ if (PyStackRef_IsTaggedInt(stackref)) { \
+ return false; \
+ } \
+ return Py ## T ## _Check(PyStackRef_AsPyObjectBorrow(stackref)); \
+ }
 
-static inline bool
-PyStackRef_BoolCheck(_PyStackRef stackref)
-{
- return PyBool_Check(PyStackRef_AsPyObjectBorrow(stackref));
-}
+STACKREF_CHECK_FUNC(Gen)
+STACKREF_CHECK_FUNC(Bool)
+STACKREF_CHECK_FUNC(ExceptionInstance)
+STACKREF_CHECK_FUNC(Code)
+STACKREF_CHECK_FUNC(Function)
 
 static inline bool
 PyStackRef_LongCheck(_PyStackRef stackref)
 {
+ if (PyStackRef_IsTaggedInt(stackref)) {
+ return true;
+ }
  return PyLong_Check(PyStackRef_AsPyObjectBorrow(stackref));
 }
 
-static inline bool
-PyStackRef_ExceptionInstanceCheck(_PyStackRef stackref)
-{
- return PyExceptionInstance_Check(PyStackRef_AsPyObjectBorrow(stackref));
-}
-
-static inline bool
-PyStackRef_CodeCheck(_PyStackRef stackref)
-{
- return PyCode_Check(PyStackRef_AsPyObjectBorrow(stackref));
-}
-
-static inline bool
-PyStackRef_FunctionCheck(_PyStackRef stackref)
-{
- return PyFunction_Check(PyStackRef_AsPyObjectBorrow(stackref));
-}
-
 static inline void
 _PyThreadState_PushCStackRef(PyThreadState *tstate, _PyCStackRef *ref)
 {