GH-96793: Change FOR_ITER to not pop the iterator on exhaustion.

Doc/library/dis.rst

@@ -1072,9 +1072,11 @@ iterations of the loop.
 
    TOS is an :term:`iterator`.  Call its :meth:`~iterator.__next__` method.  If
    this yields a new value, push it on the stack (leaving the iterator below
-   it).  If the iterator indicates it is exhausted, TOS is popped, and the byte
+   it).  If the iterator indicates it is exhausted then the byte
    code counter is incremented by *delta*.
 
+   .. versionchanged:: 3.12
+      Up until 3.11 the iterator was popped when it was exhausted.
 
 .. opcode:: LOAD_GLOBAL (namei)
 

Include/internal/pycore_code.h

@@ -250,7 +250,7 @@ extern void _Py_Specialize_CompareOp(PyObject *lhs, PyObject *rhs,
                                      _Py_CODEUNIT *instr, int oparg);
 extern void _Py_Specialize_UnpackSequence(PyObject *seq, _Py_CODEUNIT *instr,
                                           int oparg);
-extern void _Py_Specialize_ForIter(PyObject *iter, _Py_CODEUNIT *instr);
+extern void _Py_Specialize_ForIter(PyObject *iter, _Py_CODEUNIT *instr, int oparg);
 
 /* Deallocator function for static codeobjects used in deepfreeze.py */
 extern void _PyStaticCode_Dealloc(PyCodeObject *co);

Lib/importlib/_bootstrap_external.py

@@ -423,6 +423,7 @@ def _write_atomic(path, data, mode=0o666):
 #     Python 3.12a1 3507 (Set lineno of module's RESUME to 0)
 #     Python 3.12a1 3508 (Add CLEANUP_THROW)
 #     Python 3.12a1 3509 (Conditional jumps only jump forward)
+#     Python 3.12a1 3510 (FOR_ITER leaves iterator on the stack)
 
 #     Python 3.13 will start with 3550
 
@@ -435,7 +436,7 @@ def _write_atomic(path, data, mode=0o666):
 # Whenever MAGIC_NUMBER is changed, the ranges in the magic_values array
 # in PC/launcher.c must also be updated.
 
-MAGIC_NUMBER = (3509).to_bytes(2, 'little') + b'\r\n'
+MAGIC_NUMBER = (3510).to_bytes(2, 'little') + b'\r\n'
 
 _RAW_MAGIC_NUMBER = int.from_bytes(MAGIC_NUMBER, 'little')  # For import.c
 

Lib/test/test__opcode.py

@@ -40,7 +40,7 @@ def test_stack_effect_jump(self):
         self.assertEqual(stack_effect(JUMP_IF_TRUE_OR_POP, 0, jump=False), -1)
         FOR_ITER = dis.opmap['FOR_ITER']
         self.assertEqual(stack_effect(FOR_ITER, 0), 1)
-        self.assertEqual(stack_effect(FOR_ITER, 0, jump=True), -1)
+        self.assertEqual(stack_effect(FOR_ITER, 0, jump=True), 0)
         self.assertEqual(stack_effect(FOR_ITER, 0, jump=False), 1)
         JUMP_FORWARD = dis.opmap['JUMP_FORWARD']
         self.assertEqual(stack_effect(JUMP_FORWARD, 0), 0)

Misc/NEWS.d/next/Core and Builtins/2022-09-13-14-07-06.gh-issue-96793.7DLRSm.rst

@@ -0,0 +1,2 @@
+The :opcode:`FOR_ITER` now leaves the iterator on the stack on termination
+of the loop. This is to assist specialization of loops for generators.

Objects/frameobject.c

@@ -367,7 +367,7 @@ mark_stacks(PyCodeObject *code_obj, int len)
                     break;
                 case FOR_ITER:
                 {
-                    int64_t target_stack = pop_value(next_stack);
+                    int64_t target_stack = next_stack;
                     stacks[i+1] = push_value(next_stack, Object);
                     j = get_arg(code, i) + 1 + INLINE_CACHE_ENTRIES_FOR_ITER + i;
                     assert(j < len);

Python/ceval.c

@@ -3844,8 +3844,6 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
                 _PyErr_Clear(tstate);
             }
             /* iterator ended normally */
-            STACK_SHRINK(1);
-            Py_DECREF(iter);
             JUMPBY(INLINE_CACHE_ENTRIES_FOR_ITER + oparg);
             DISPATCH();
         }
@@ -3855,7 +3853,7 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
             _PyForIterCache *cache = (_PyForIterCache *)next_instr;
             if (ADAPTIVE_COUNTER_IS_ZERO(cache)) {
                 next_instr--;
-                _Py_Specialize_ForIter(TOP(), next_instr);
+                _Py_Specialize_ForIter(TOP(), next_instr, oparg);
                 DISPATCH_SAME_OPARG();
             }
             else {
@@ -3884,7 +3882,7 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
             }
             STACK_SHRINK(1);
             Py_DECREF(it);
-            JUMPBY(INLINE_CACHE_ENTRIES_FOR_ITER + oparg);
+            JUMPBY(INLINE_CACHE_ENTRIES_FOR_ITER + oparg + 1);
             DISPATCH();
         }
 
@@ -3898,7 +3896,7 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
             if (r->index >= r->len) {
                 STACK_SHRINK(1);
                 Py_DECREF(r);
-                JUMPBY(INLINE_CACHE_ENTRIES_FOR_ITER + oparg);
+                JUMPBY(INLINE_CACHE_ENTRIES_FOR_ITER + oparg + 1);
                 DISPATCH();
             }
             long value = (long)(r->start +

Python/compile.c

@@ -1099,7 +1099,7 @@ stack_effect(int opcode, int oparg, int jump)
             return (oparg&0xFF) + (oparg>>8);
         case FOR_ITER:
             /* -1 at end of iterator, 1 if continue iterating. */
-            return jump > 0 ? -1 : 1;
+            return jump > 0 ? 0 : 1;
         case SEND:
             return jump > 0 ? -1 : 0;
         case STORE_ATTR:
@@ -3103,6 +3103,7 @@ compiler_for(struct compiler *c, stmt_ty s)
     ADDOP_JUMP(c, NO_LOCATION, JUMP, start);
 
     USE_LABEL(c, cleanup);
+    ADDOP(c, NO_LOCATION, POP_TOP);
 
     compiler_pop_fblock(c, FOR_LOOP, start);
 
@@ -5315,6 +5316,7 @@ compiler_sync_comprehension_generator(struct compiler *c, location loc,
         ADDOP_JUMP(c, elt_loc, JUMP, start);
 
         USE_LABEL(c, anchor);
+        ADDOP(c, NO_LOCATION, POP_TOP);
     }
 
     return 1;

Python/specialize.c

@@ -2193,27 +2193,25 @@ int
 #endif
 
 void
-_Py_Specialize_ForIter(PyObject *iter, _Py_CODEUNIT *instr)
+_Py_Specialize_ForIter(PyObject *iter, _Py_CODEUNIT *instr, int oparg)
 {
     assert(_PyOpcode_Caches[FOR_ITER] == INLINE_CACHE_ENTRIES_FOR_ITER);
     _PyForIterCache *cache = (_PyForIterCache *)(instr + 1);
     PyTypeObject *tp = Py_TYPE(iter);
-    _Py_CODEUNIT next = instr[1+INLINE_CACHE_ENTRIES_FOR_ITER];
-    int next_op = _PyOpcode_Deopt[_Py_OPCODE(next)];
-    if (tp == &PyListIter_Type) {
-        _Py_SET_OPCODE(*instr, FOR_ITER_LIST);
-        goto success;
-    }
-    else if (tp == &PyRangeIter_Type && next_op == STORE_FAST) {
-        _Py_SET_OPCODE(*instr, FOR_ITER_RANGE);
-        goto success;
-    }
-    else {
-        SPECIALIZATION_FAIL(FOR_ITER,
-                            _PySpecialization_ClassifyIterator(iter));
-        goto failure;
+    if (_Py_OPCODE(instr[INLINE_CACHE_ENTRIES_FOR_ITER + oparg + 1]) == POP_TOP) {
+        _Py_CODEUNIT next = instr[1+INLINE_CACHE_ENTRIES_FOR_ITER];
+        int next_op = _PyOpcode_Deopt[_Py_OPCODE(next)];
+        if (tp == &PyListIter_Type) {
+            _Py_SET_OPCODE(*instr, FOR_ITER_LIST);
+            goto success;
+        }
+        else if (tp == &PyRangeIter_Type && next_op == STORE_FAST) {
+            _Py_SET_OPCODE(*instr, FOR_ITER_RANGE);
+            goto success;
+        }
     }
-failure:
+    SPECIALIZATION_FAIL(FOR_ITER,
+                        _PySpecialization_ClassifyIterator(iter));
     STAT_INC(FOR_ITER, failure);
     cache->counter = adaptive_counter_backoff(cache->counter);
     return;