GH-96793: Change `FOR_ITER` to not pop the iterator on exhaustion. (GH-96801)

Change FOR_ITER to have the same stack effect regardless of whether it branches or not.
Performance is unchanged as FOR_ITER (and specialized forms jump over the cleanup code).

4 files changed, 39 insertions, 30 deletions

diff --git a/Python/ceval.c b/Python/ceval.c
index 35ce767710e..ac995d7c178 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -1298,6 +1298,14 @@ handle_eval_breaker:
             DISPATCH();
         }
 
+        TARGET(END_FOR) {
+            PyObject *value = POP();
+            Py_DECREF(value);
+            value = POP();
+            Py_DECREF(value);
+            DISPATCH();
+        }
+
         TARGET(UNARY_POSITIVE) {
             PyObject *value = TOP();
             PyObject *res = PyNumber_Positive(value);
@@ -3844,9 +3852,11 @@ handle_eval_breaker:
                 _PyErr_Clear(tstate);
             }
             /* iterator ended normally */
+            assert(_Py_OPCODE(next_instr[INLINE_CACHE_ENTRIES_FOR_ITER + oparg] == END_FOR));
             STACK_SHRINK(1);
             Py_DECREF(iter);
-            JUMPBY(INLINE_CACHE_ENTRIES_FOR_ITER + oparg);
+            /* Skip END_FOR */
+            JUMPBY(INLINE_CACHE_ENTRIES_FOR_ITER + oparg + 1);
             DISPATCH();
         }
 
@@ -3884,7 +3894,7 @@ handle_eval_breaker:
             }
             STACK_SHRINK(1);
             Py_DECREF(it);
-            JUMPBY(INLINE_CACHE_ENTRIES_FOR_ITER + oparg);
+            JUMPBY(INLINE_CACHE_ENTRIES_FOR_ITER + oparg + 1);
             DISPATCH();
         }
 
@@ -3898,7 +3908,7 @@ handle_eval_breaker:
             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 +
diff --git a/Python/compile.c b/Python/compile.c
index 67ccd223938..779729119e2 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -1042,6 +1042,8 @@ stack_effect(int opcode, int oparg, int jump)
             return -1;
         case SWAP:
             return 0;
+        case END_FOR:
+            return -2;
 
         /* Unary operators */
         case UNARY_POSITIVE:
@@ -1098,8 +1100,7 @@ stack_effect(int opcode, int oparg, int jump)
         case UNPACK_EX:
             return (oparg&0xFF) + (oparg>>8);
         case FOR_ITER:
-            /* -1 at end of iterator, 1 if continue iterating. */
-            return jump > 0 ? -1 : 1;
+            return 1;
         case SEND:
             return jump > 0 ? -1 : 0;
         case STORE_ATTR:
@@ -3103,6 +3104,7 @@ compiler_for(struct compiler *c, stmt_ty s)
     ADDOP_JUMP(c, NO_LOCATION, JUMP, start);
 
     USE_LABEL(c, cleanup);
+    ADDOP(c, NO_LOCATION, END_FOR);
 
     compiler_pop_fblock(c, FOR_LOOP, start);
 
@@ -5315,6 +5317,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, END_FOR);
     }
 
     return 1;
diff --git a/Python/opcode_targets.h b/Python/opcode_targets.h
index c1ff367d4fd..53b2770607b 100644
--- a/Python/opcode_targets.h
+++ b/Python/opcode_targets.h
@@ -3,18 +3,19 @@ static void *opcode_targets[256] = {
     &&TARGET_POP_TOP,
     &&TARGET_PUSH_NULL,
     &&TARGET_BINARY_OP_ADAPTIVE,
+    &&TARGET_END_FOR,
     &&TARGET_BINARY_OP_ADD_FLOAT,
     &&TARGET_BINARY_OP_ADD_INT,
     &&TARGET_BINARY_OP_ADD_UNICODE,
     &&TARGET_BINARY_OP_INPLACE_ADD_UNICODE,
-    &&TARGET_BINARY_OP_MULTIPLY_FLOAT,
     &&TARGET_NOP,
     &&TARGET_UNARY_POSITIVE,
     &&TARGET_UNARY_NEGATIVE,
     &&TARGET_UNARY_NOT,
+    &&TARGET_BINARY_OP_MULTIPLY_FLOAT,
     &&TARGET_BINARY_OP_MULTIPLY_INT,
-    &&TARGET_BINARY_OP_SUBTRACT_FLOAT,
     &&TARGET_UNARY_INVERT,
+    &&TARGET_BINARY_OP_SUBTRACT_FLOAT,
     &&TARGET_BINARY_OP_SUBTRACT_INT,
     &&TARGET_BINARY_SUBSCR_ADAPTIVE,
     &&TARGET_BINARY_SUBSCR_DICT,
@@ -23,20 +24,20 @@ static void *opcode_targets[256] = {
     &&TARGET_BINARY_SUBSCR_TUPLE_INT,
     &&TARGET_CALL_ADAPTIVE,
     &&TARGET_CALL_PY_EXACT_ARGS,
-    &&TARGET_CALL_PY_WITH_DEFAULTS,
     &&TARGET_BINARY_SUBSCR,
     &&TARGET_BINARY_SLICE,
     &&TARGET_STORE_SLICE,
+    &&TARGET_CALL_PY_WITH_DEFAULTS,
     &&TARGET_CALL_BOUND_METHOD_EXACT_ARGS,
-    &&TARGET_CALL_BUILTIN_CLASS,
     &&TARGET_GET_LEN,
     &&TARGET_MATCH_MAPPING,
     &&TARGET_MATCH_SEQUENCE,
     &&TARGET_MATCH_KEYS,
-    &&TARGET_CALL_BUILTIN_FAST_WITH_KEYWORDS,
+    &&TARGET_CALL_BUILTIN_CLASS,
     &&TARGET_PUSH_EXC_INFO,
     &&TARGET_CHECK_EXC_MATCH,
     &&TARGET_CHECK_EG_MATCH,
+    &&TARGET_CALL_BUILTIN_FAST_WITH_KEYWORDS,
     &&TARGET_CALL_METHOD_DESCRIPTOR_FAST_WITH_KEYWORDS,
     &&TARGET_CALL_NO_KW_BUILTIN_FAST,
     &&TARGET_CALL_NO_KW_BUILTIN_O,
@@ -47,7 +48,6 @@ static void *opcode_targets[256] = {
     &&TARGET_CALL_NO_KW_METHOD_DESCRIPTOR_NOARGS,
     &&TARGET_CALL_NO_KW_METHOD_DESCRIPTOR_O,
     &&TARGET_CALL_NO_KW_STR_1,
-    &&TARGET_CALL_NO_KW_TUPLE_1,
     &&TARGET_WITH_EXCEPT_START,
     &&TARGET_GET_AITER,
     &&TARGET_GET_ANEXT,
@@ -55,37 +55,37 @@ static void *opcode_targets[256] = {
     &&TARGET_BEFORE_WITH,
     &&TARGET_END_ASYNC_FOR,
     &&TARGET_CLEANUP_THROW,
+    &&TARGET_CALL_NO_KW_TUPLE_1,
     &&TARGET_CALL_NO_KW_TYPE_1,
     &&TARGET_COMPARE_OP_ADAPTIVE,
     &&TARGET_COMPARE_OP_FLOAT_JUMP,
-    &&TARGET_COMPARE_OP_INT_JUMP,
     &&TARGET_STORE_SUBSCR,
     &&TARGET_DELETE_SUBSCR,
+    &&TARGET_COMPARE_OP_INT_JUMP,
     &&TARGET_COMPARE_OP_STR_JUMP,
     &&TARGET_EXTENDED_ARG_QUICK,
     &&TARGET_FOR_ITER_ADAPTIVE,
     &&TARGET_FOR_ITER_LIST,
     &&TARGET_FOR_ITER_RANGE,
-    &&TARGET_JUMP_BACKWARD_QUICK,
     &&TARGET_GET_ITER,
     &&TARGET_GET_YIELD_FROM_ITER,
     &&TARGET_PRINT_EXPR,
     &&TARGET_LOAD_BUILD_CLASS,
+    &&TARGET_JUMP_BACKWARD_QUICK,
     &&TARGET_LOAD_ATTR_ADAPTIVE,
-    &&TARGET_LOAD_ATTR_CLASS,
     &&TARGET_LOAD_ASSERTION_ERROR,
     &&TARGET_RETURN_GENERATOR,
+    &&TARGET_LOAD_ATTR_CLASS,
     &&TARGET_LOAD_ATTR_GETATTRIBUTE_OVERRIDDEN,
     &&TARGET_LOAD_ATTR_INSTANCE_VALUE,
     &&TARGET_LOAD_ATTR_MODULE,
     &&TARGET_LOAD_ATTR_PROPERTY,
     &&TARGET_LOAD_ATTR_SLOT,
-    &&TARGET_LOAD_ATTR_WITH_HINT,
     &&TARGET_LIST_TO_TUPLE,
     &&TARGET_RETURN_VALUE,
     &&TARGET_IMPORT_STAR,
     &&TARGET_SETUP_ANNOTATIONS,
-    &&TARGET_LOAD_ATTR_METHOD_LAZY_DICT,
+    &&TARGET_LOAD_ATTR_WITH_HINT,
     &&TARGET_ASYNC_GEN_WRAP,
     &&TARGET_PREP_RERAISE_STAR,
     &&TARGET_POP_EXCEPT,
@@ -112,7 +112,7 @@ static void *opcode_targets[256] = {
     &&TARGET_JUMP_FORWARD,
     &&TARGET_JUMP_IF_FALSE_OR_POP,
     &&TARGET_JUMP_IF_TRUE_OR_POP,
-    &&TARGET_LOAD_ATTR_METHOD_NO_DICT,
+    &&TARGET_LOAD_ATTR_METHOD_LAZY_DICT,
     &&TARGET_POP_JUMP_IF_FALSE,
     &&TARGET_POP_JUMP_IF_TRUE,
     &&TARGET_LOAD_GLOBAL,
@@ -120,7 +120,7 @@ static void *opcode_targets[256] = {
     &&TARGET_CONTAINS_OP,
     &&TARGET_RERAISE,
     &&TARGET_COPY,
-    &&TARGET_LOAD_ATTR_METHOD_WITH_DICT,
+    &&TARGET_LOAD_ATTR_METHOD_NO_DICT,
     &&TARGET_BINARY_OP,
     &&TARGET_SEND,
     &&TARGET_LOAD_FAST,
@@ -140,9 +140,9 @@ static void *opcode_targets[256] = {
     &&TARGET_STORE_DEREF,
     &&TARGET_DELETE_DEREF,
     &&TARGET_JUMP_BACKWARD,
-    &&TARGET_LOAD_ATTR_METHOD_WITH_VALUES,
+    &&TARGET_LOAD_ATTR_METHOD_WITH_DICT,
     &&TARGET_CALL_FUNCTION_EX,
-    &&TARGET_LOAD_CONST__LOAD_FAST,
+    &&TARGET_LOAD_ATTR_METHOD_WITH_VALUES,
     &&TARGET_EXTENDED_ARG,
     &&TARGET_LIST_APPEND,
     &&TARGET_SET_ADD,
@@ -152,26 +152,27 @@ static void *opcode_targets[256] = {
     &&TARGET_YIELD_VALUE,
     &&TARGET_RESUME,
     &&TARGET_MATCH_CLASS,
+    &&TARGET_LOAD_CONST__LOAD_FAST,
     &&TARGET_LOAD_FAST__LOAD_CONST,
-    &&TARGET_LOAD_FAST__LOAD_FAST,
     &&TARGET_FORMAT_VALUE,
     &&TARGET_BUILD_CONST_KEY_MAP,
     &&TARGET_BUILD_STRING,
+    &&TARGET_LOAD_FAST__LOAD_FAST,
     &&TARGET_LOAD_GLOBAL_ADAPTIVE,
     &&TARGET_LOAD_GLOBAL_BUILTIN,
     &&TARGET_LOAD_GLOBAL_MODULE,
-    &&TARGET_RESUME_QUICK,
     &&TARGET_LIST_EXTEND,
     &&TARGET_SET_UPDATE,
     &&TARGET_DICT_MERGE,
     &&TARGET_DICT_UPDATE,
+    &&TARGET_RESUME_QUICK,
     &&TARGET_STORE_ATTR_ADAPTIVE,
     &&TARGET_STORE_ATTR_INSTANCE_VALUE,
     &&TARGET_STORE_ATTR_SLOT,
     &&TARGET_STORE_ATTR_WITH_HINT,
-    &&TARGET_STORE_FAST__LOAD_FAST,
     &&TARGET_CALL,
     &&TARGET_KW_NAMES,
+    &&TARGET_STORE_FAST__LOAD_FAST,
     &&TARGET_STORE_FAST__STORE_FAST,
     &&TARGET_STORE_SUBSCR_ADAPTIVE,
     &&TARGET_STORE_SUBSCR_DICT,
@@ -253,6 +254,5 @@ static void *opcode_targets[256] = {
     &&_unknown_opcode,
     &&_unknown_opcode,
     &&_unknown_opcode,
-    &&_unknown_opcode,
     &&TARGET_DO_TRACING
 };
diff --git a/Python/specialize.c b/Python/specialize.c
index b7c321e4878..6e2fa4e25fa 100644
--- a/Python/specialize.c
+++ b/Python/specialize.c
@@ -2208,12 +2208,8 @@ _Py_Specialize_ForIter(PyObject *iter, _Py_CODEUNIT *instr)
         _Py_SET_OPCODE(*instr, FOR_ITER_RANGE);
         goto success;
     }
-    else {
-        SPECIALIZATION_FAIL(FOR_ITER,
-                            _PySpecialization_ClassifyIterator(iter));
-        goto failure;
-    }
-failure:
+    SPECIALIZATION_FAIL(FOR_ITER,
+                        _PySpecialization_ClassifyIterator(iter));
     STAT_INC(FOR_ITER, failure);
     cache->counter = adaptive_counter_backoff(cache->counter);
     return;