GH-96793: Specialize FOR_ITER for generators. (GH-98772)

5 files changed, 132 insertions, 23 deletions

diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index f0e9e3a..636d2d8 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -756,6 +756,11 @@ dummy_func(
                 goto resume_frame;
             }
             _Py_LeaveRecursiveCallTstate(tstate);
+            if (frame->owner == FRAME_OWNED_BY_GENERATOR) {
+                PyGenObject *gen = _PyFrame_GetGenerator(frame);
+                tstate->exc_info = gen->gi_exc_state.previous_item;
+                gen->gi_exc_state.previous_item = NULL;
+            }
             /* Restore previous cframe and return. */
             tstate->cframe = cframe.previous;
             tstate->cframe->use_tracing = cframe.use_tracing;
@@ -895,7 +900,6 @@ dummy_func(
 
         // error: SEND stack effect depends on jump flag
         inst(SEND) {
-            assert(frame->is_entry);
             assert(STACK_LEVEL() >= 2);
             PyObject *v = POP();
             PyObject *receiver = TOP();
@@ -960,13 +964,21 @@ dummy_func(
             // The compiler treats any exception raised here as a failed close()
             // or throw() call.
             assert(oparg == STACK_LEVEL());
-            assert(frame->is_entry);
             PyObject *retval = POP();
-            _PyFrame_GetGenerator(frame)->gi_frame_state = FRAME_SUSPENDED;
+            PyGenObject *gen = _PyFrame_GetGenerator(frame);
+            gen->gi_frame_state = FRAME_SUSPENDED;
             _PyFrame_SetStackPointer(frame, stack_pointer);
             TRACE_FUNCTION_EXIT();
             DTRACE_FUNCTION_EXIT();
+            tstate->exc_info = gen->gi_exc_state.previous_item;
+            gen->gi_exc_state.previous_item = NULL;
             _Py_LeaveRecursiveCallPy(tstate);
+            if (!frame->is_entry) {
+                frame = cframe.current_frame = frame->previous;
+                frame->prev_instr -= frame->yield_offset;
+                _PyFrame_StackPush(frame, retval);
+                goto resume_frame;
+            }
             _Py_LeaveRecursiveCallTstate(tstate);
             /* Restore previous cframe and return. */
             tstate->cframe = cframe.previous;
@@ -2788,7 +2800,7 @@ dummy_func(
             _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 {
@@ -2844,6 +2856,30 @@ dummy_func(
             JUMPBY(INLINE_CACHE_ENTRIES_FOR_ITER + 1);
         }
 
+        inst(FOR_ITER_GEN) {
+            assert(cframe.use_tracing == 0);
+            PyGenObject *gen = (PyGenObject *)TOP();
+            DEOPT_IF(Py_TYPE(gen) != &PyGen_Type, FOR_ITER);
+            DEOPT_IF(gen->gi_frame_state >= FRAME_EXECUTING, FOR_ITER);
+            STAT_INC(FOR_ITER, hit);
+            _PyInterpreterFrame *gen_frame = (_PyInterpreterFrame *)gen->gi_iframe;
+            _PyFrame_SetStackPointer(frame, stack_pointer);
+            frame->yield_offset = oparg;
+            JUMPBY(INLINE_CACHE_ENTRIES_FOR_ITER + oparg);
+            assert(_Py_OPCODE(*next_instr) == END_FOR);
+            frame->prev_instr = next_instr - 1;
+            Py_INCREF(Py_None);
+            _PyFrame_StackPush(gen_frame, Py_None);
+            gen->gi_frame_state = FRAME_EXECUTING;
+            gen->gi_exc_state.previous_item = tstate->exc_info;
+            tstate->exc_info = &gen->gi_exc_state;
+            gen_frame->previous = frame;
+            gen_frame->is_entry = false;
+            frame = cframe.current_frame = gen_frame;
+            goto start_frame;
+        }
+
+
         // stack effect: ( -- __0)
         inst(BEFORE_ASYNC_WITH) {
             PyObject *mgr = TOP();
diff --git a/Python/ceval.c b/Python/ceval.c
index 85011af..46fad97 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -1360,6 +1360,11 @@ exit_unwind:
     assert(_PyErr_Occurred(tstate));
     _Py_LeaveRecursiveCallPy(tstate);
     if (frame->is_entry) {
+        if (frame->owner == FRAME_OWNED_BY_GENERATOR) {
+            PyGenObject *gen = _PyFrame_GetGenerator(frame);
+            tstate->exc_info = gen->gi_exc_state.previous_item;
+            gen->gi_exc_state.previous_item = NULL;
+        }
         /* Restore previous cframe and exit */
         tstate->cframe = cframe.previous;
         tstate->cframe->use_tracing = cframe.use_tracing;
@@ -1930,20 +1935,48 @@ fail:
 }
 
 static void
-_PyEvalFrameClearAndPop(PyThreadState *tstate, _PyInterpreterFrame * frame)
+clear_thread_frame(PyThreadState *tstate, _PyInterpreterFrame * frame)
 {
+    assert(frame->owner == FRAME_OWNED_BY_THREAD);
     // Make sure that this is, indeed, the top frame. We can't check this in
     // _PyThreadState_PopFrame, since f_code is already cleared at that point:
     assert((PyObject **)frame + frame->f_code->co_framesize ==
-           tstate->datastack_top);
+        tstate->datastack_top);
     tstate->c_recursion_remaining--;
     assert(frame->frame_obj == NULL || frame->frame_obj->f_frame == frame);
-    assert(frame->owner == FRAME_OWNED_BY_THREAD);
     _PyFrame_Clear(frame);
     tstate->c_recursion_remaining++;
     _PyThreadState_PopFrame(tstate, frame);
 }
 
+static void
+clear_gen_frame(PyThreadState *tstate, _PyInterpreterFrame * frame)
+{
+    assert(frame->owner == FRAME_OWNED_BY_GENERATOR);
+    PyGenObject *gen = _PyFrame_GetGenerator(frame);
+    gen->gi_frame_state = FRAME_CLEARED;
+    assert(tstate->exc_info == &gen->gi_exc_state);
+    tstate->exc_info = gen->gi_exc_state.previous_item;
+    gen->gi_exc_state.previous_item = NULL;
+    tstate->c_recursion_remaining--;
+    assert(frame->frame_obj == NULL || frame->frame_obj->f_frame == frame);
+    _PyFrame_Clear(frame);
+    tstate->c_recursion_remaining++;
+    frame->previous = NULL;
+}
+
+static void
+_PyEvalFrameClearAndPop(PyThreadState *tstate, _PyInterpreterFrame * frame)
+{
+    if (frame->owner == FRAME_OWNED_BY_THREAD) {
+        clear_thread_frame(tstate, frame);
+    }
+    else {
+        clear_gen_frame(tstate, frame);
+    }
+}
+
+
 PyObject *
 _PyEval_Vector(PyThreadState *tstate, PyFunctionObject *func,
                PyObject *locals,
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
index d83d683..6ea1f48 100644
--- a/Python/generated_cases.c.h
+++ b/Python/generated_cases.c.h
@@ -646,6 +646,11 @@
                 goto resume_frame;
             }
             _Py_LeaveRecursiveCallTstate(tstate);
+            if (frame->owner == FRAME_OWNED_BY_GENERATOR) {
+                PyGenObject *gen = _PyFrame_GetGenerator(frame);
+                tstate->exc_info = gen->gi_exc_state.previous_item;
+                gen->gi_exc_state.previous_item = NULL;
+            }
             /* Restore previous cframe and return. */
             tstate->cframe = cframe.previous;
             tstate->cframe->use_tracing = cframe.use_tracing;
@@ -785,7 +790,6 @@
         }
 
         TARGET(SEND) {
-            assert(frame->is_entry);
             assert(STACK_LEVEL() >= 2);
             PyObject *v = POP();
             PyObject *receiver = TOP();
@@ -850,13 +854,21 @@
             // The compiler treats any exception raised here as a failed close()
             // or throw() call.
             assert(oparg == STACK_LEVEL());
-            assert(frame->is_entry);
             PyObject *retval = POP();
-            _PyFrame_GetGenerator(frame)->gi_frame_state = FRAME_SUSPENDED;
+            PyGenObject *gen = _PyFrame_GetGenerator(frame);
+            gen->gi_frame_state = FRAME_SUSPENDED;
             _PyFrame_SetStackPointer(frame, stack_pointer);
             TRACE_FUNCTION_EXIT();
             DTRACE_FUNCTION_EXIT();
+            tstate->exc_info = gen->gi_exc_state.previous_item;
+            gen->gi_exc_state.previous_item = NULL;
             _Py_LeaveRecursiveCallPy(tstate);
+            if (!frame->is_entry) {
+                frame = cframe.current_frame = frame->previous;
+                frame->prev_instr -= frame->yield_offset;
+                _PyFrame_StackPush(frame, retval);
+                goto resume_frame;
+            }
             _Py_LeaveRecursiveCallTstate(tstate);
             /* Restore previous cframe and return. */
             tstate->cframe = cframe.previous;
@@ -2678,7 +2690,7 @@
             _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 {
@@ -2735,6 +2747,29 @@
             DISPATCH();
         }
 
+        TARGET(FOR_ITER_GEN) {
+            assert(cframe.use_tracing == 0);
+            PyGenObject *gen = (PyGenObject *)TOP();
+            DEOPT_IF(Py_TYPE(gen) != &PyGen_Type, FOR_ITER);
+            DEOPT_IF(gen->gi_frame_state >= FRAME_EXECUTING, FOR_ITER);
+            STAT_INC(FOR_ITER, hit);
+            _PyInterpreterFrame *gen_frame = (_PyInterpreterFrame *)gen->gi_iframe;
+            _PyFrame_SetStackPointer(frame, stack_pointer);
+            frame->yield_offset = oparg;
+            JUMPBY(INLINE_CACHE_ENTRIES_FOR_ITER + oparg);
+            assert(_Py_OPCODE(*next_instr) == END_FOR);
+            frame->prev_instr = next_instr - 1;
+            Py_INCREF(Py_None);
+            _PyFrame_StackPush(gen_frame, Py_None);
+            gen->gi_frame_state = FRAME_EXECUTING;
+            gen->gi_exc_state.previous_item = tstate->exc_info;
+            tstate->exc_info = &gen->gi_exc_state;
+            gen_frame->previous = frame;
+            gen_frame->is_entry = false;
+            frame = cframe.current_frame = gen_frame;
+            goto start_frame;
+        }
+
         TARGET(BEFORE_ASYNC_WITH) {
             PyObject *mgr = TOP();
             PyObject *res;
diff --git a/Python/opcode_targets.h b/Python/opcode_targets.h
index a963a7a..2251c94 100644
--- a/Python/opcode_targets.h
+++ b/Python/opcode_targets.h
@@ -72,20 +72,20 @@ static void *opcode_targets[256] = {
     &&TARGET_PRINT_EXPR,
     &&TARGET_LOAD_BUILD_CLASS,
     &&TARGET_FOR_ITER_RANGE,
-    &&TARGET_LOAD_ATTR_ADAPTIVE,
+    &&TARGET_FOR_ITER_GEN,
     &&TARGET_LOAD_ASSERTION_ERROR,
     &&TARGET_RETURN_GENERATOR,
+    &&TARGET_LOAD_ATTR_ADAPTIVE,
     &&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_LIST_TO_TUPLE,
     &&TARGET_RETURN_VALUE,
     &&TARGET_IMPORT_STAR,
     &&TARGET_SETUP_ANNOTATIONS,
-    &&TARGET_LOAD_ATTR_WITH_HINT,
+    &&TARGET_LOAD_ATTR_SLOT,
     &&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_LAZY_DICT,
+    &&TARGET_LOAD_ATTR_WITH_HINT,
     &&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_NO_DICT,
+    &&TARGET_LOAD_ATTR_METHOD_LAZY_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_DICT,
+    &&TARGET_LOAD_ATTR_METHOD_NO_DICT,
     &&TARGET_CALL_FUNCTION_EX,
-    &&TARGET_LOAD_ATTR_METHOD_WITH_VALUES,
+    &&TARGET_LOAD_ATTR_METHOD_WITH_DICT,
     &&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_ATTR_METHOD_WITH_VALUES,
     &&TARGET_LOAD_CONST__LOAD_FAST,
-    &&TARGET_LOAD_FAST__LOAD_CONST,
     &&TARGET_FORMAT_VALUE,
     &&TARGET_BUILD_CONST_KEY_MAP,
     &&TARGET_BUILD_STRING,
+    &&TARGET_LOAD_FAST__LOAD_CONST,
     &&TARGET_LOAD_FAST__LOAD_FAST,
     &&TARGET_LOAD_GLOBAL_ADAPTIVE,
     &&TARGET_LOAD_GLOBAL_BUILTIN,
-    &&TARGET_LOAD_GLOBAL_MODULE,
     &&TARGET_LIST_EXTEND,
     &&TARGET_SET_UPDATE,
     &&TARGET_DICT_MERGE,
     &&TARGET_DICT_UPDATE,
+    &&TARGET_LOAD_GLOBAL_MODULE,
     &&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 70a456c..5717991 100644
--- a/Python/specialize.c
+++ b/Python/specialize.c
@@ -2184,7 +2184,7 @@ 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);
@@ -2199,6 +2199,11 @@ _Py_Specialize_ForIter(PyObject *iter, _Py_CODEUNIT *instr)
         _Py_SET_OPCODE(*instr, FOR_ITER_RANGE);
         goto success;
     }
+    else if (tp == &PyGen_Type && oparg <= SHRT_MAX) {
+        assert(_Py_OPCODE(instr[oparg + INLINE_CACHE_ENTRIES_FOR_ITER + 1]) == END_FOR);
+        _Py_SET_OPCODE(*instr, FOR_ITER_GEN);
+        goto success;
+    }
     SPECIALIZATION_FAIL(FOR_ITER,
                         _PySpecialization_ClassifyIterator(iter));
     STAT_INC(FOR_ITER, failure);