pythonGH-132554: "Virtual" iterators (pythonGH-132555)

* FOR_ITER now pushes either the iterator and NULL or leaves the iterable and pushes tagged zero

* NEXT_ITER uses the tagged int as the index into the sequence or, if TOS is NULL, iterates as before.

Author: markshannon

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_opcode_metadata.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)
 {