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

Include/internal/pycore_code.h

@@ -230,7 +230,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);
 
 /* Finalizer function for static codeobjects used in deepfreeze.py */
 extern void _PyStaticCode_Fini(PyCodeObject *co);

Include/internal/pycore_frame.h

@@ -61,6 +61,7 @@ typedef struct _PyInterpreterFrame {
     // over, or (in the case of a newly-created frame) a totally invalid value:
     _Py_CODEUNIT *prev_instr;
     int stacktop;     /* Offset of TOS from localsplus  */
+    uint16_t yield_offset;
     bool is_entry;  // Whether this is the "root" frame for the current _PyCFrame.
     char owner;
     /* Locals and stack */
@@ -110,6 +111,7 @@ _PyFrame_InitializeSpecials(
     frame->frame_obj = NULL;
     frame->prev_instr = _PyCode_CODE(code) - 1;
     frame->is_entry = false;
+    frame->yield_offset = 0;
     frame->owner = FRAME_OWNED_BY_THREAD;
 }
 

Include/internal/pycore_opcode.h

@@ -128,6 +128,7 @@ const uint8_t _PyOpcode_Deopt[256] = {
     [FORMAT_VALUE] = FORMAT_VALUE,
     [FOR_ITER] = FOR_ITER,
     [FOR_ITER_ADAPTIVE] = FOR_ITER,
+    [FOR_ITER_GEN] = FOR_ITER,
     [FOR_ITER_LIST] = FOR_ITER,
     [FOR_ITER_RANGE] = FOR_ITER,
     [GET_AITER] = GET_AITER,
@@ -313,20 +314,20 @@ static const char *const _PyOpcode_OpName[263] = {
     [PRINT_EXPR] = "PRINT_EXPR",
     [LOAD_BUILD_CLASS] = "LOAD_BUILD_CLASS",
     [FOR_ITER_RANGE] = "FOR_ITER_RANGE",
-    [LOAD_ATTR_ADAPTIVE] = "LOAD_ATTR_ADAPTIVE",
+    [FOR_ITER_GEN] = "FOR_ITER_GEN",
     [LOAD_ASSERTION_ERROR] = "LOAD_ASSERTION_ERROR",
     [RETURN_GENERATOR] = "RETURN_GENERATOR",
+    [LOAD_ATTR_ADAPTIVE] = "LOAD_ATTR_ADAPTIVE",
     [LOAD_ATTR_CLASS] = "LOAD_ATTR_CLASS",
     [LOAD_ATTR_GETATTRIBUTE_OVERRIDDEN] = "LOAD_ATTR_GETATTRIBUTE_OVERRIDDEN",
     [LOAD_ATTR_INSTANCE_VALUE] = "LOAD_ATTR_INSTANCE_VALUE",
     [LOAD_ATTR_MODULE] = "LOAD_ATTR_MODULE",
     [LOAD_ATTR_PROPERTY] = "LOAD_ATTR_PROPERTY",
-    [LOAD_ATTR_SLOT] = "LOAD_ATTR_SLOT",
     [LIST_TO_TUPLE] = "LIST_TO_TUPLE",
     [RETURN_VALUE] = "RETURN_VALUE",
     [IMPORT_STAR] = "IMPORT_STAR",
     [SETUP_ANNOTATIONS] = "SETUP_ANNOTATIONS",
-    [LOAD_ATTR_WITH_HINT] = "LOAD_ATTR_WITH_HINT",
+    [LOAD_ATTR_SLOT] = "LOAD_ATTR_SLOT",
     [ASYNC_GEN_WRAP] = "ASYNC_GEN_WRAP",
     [PREP_RERAISE_STAR] = "PREP_RERAISE_STAR",
     [POP_EXCEPT] = "POP_EXCEPT",
@@ -353,7 +354,7 @@ static const char *const _PyOpcode_OpName[263] = {
     [JUMP_FORWARD] = "JUMP_FORWARD",
     [JUMP_IF_FALSE_OR_POP] = "JUMP_IF_FALSE_OR_POP",
     [JUMP_IF_TRUE_OR_POP] = "JUMP_IF_TRUE_OR_POP",
-    [LOAD_ATTR_METHOD_LAZY_DICT] = "LOAD_ATTR_METHOD_LAZY_DICT",
+    [LOAD_ATTR_WITH_HINT] = "LOAD_ATTR_WITH_HINT",
     [POP_JUMP_IF_FALSE] = "POP_JUMP_IF_FALSE",
     [POP_JUMP_IF_TRUE] = "POP_JUMP_IF_TRUE",
     [LOAD_GLOBAL] = "LOAD_GLOBAL",
@@ -361,7 +362,7 @@ static const char *const _PyOpcode_OpName[263] = {
     [CONTAINS_OP] = "CONTAINS_OP",
     [RERAISE] = "RERAISE",
     [COPY] = "COPY",
-    [LOAD_ATTR_METHOD_NO_DICT] = "LOAD_ATTR_METHOD_NO_DICT",
+    [LOAD_ATTR_METHOD_LAZY_DICT] = "LOAD_ATTR_METHOD_LAZY_DICT",
     [BINARY_OP] = "BINARY_OP",
     [SEND] = "SEND",
     [LOAD_FAST] = "LOAD_FAST",
@@ -381,9 +382,9 @@ static const char *const _PyOpcode_OpName[263] = {
     [STORE_DEREF] = "STORE_DEREF",
     [DELETE_DEREF] = "DELETE_DEREF",
     [JUMP_BACKWARD] = "JUMP_BACKWARD",
-    [LOAD_ATTR_METHOD_WITH_DICT] = "LOAD_ATTR_METHOD_WITH_DICT",
+    [LOAD_ATTR_METHOD_NO_DICT] = "LOAD_ATTR_METHOD_NO_DICT",
     [CALL_FUNCTION_EX] = "CALL_FUNCTION_EX",
-    [LOAD_ATTR_METHOD_WITH_VALUES] = "LOAD_ATTR_METHOD_WITH_VALUES",
+    [LOAD_ATTR_METHOD_WITH_DICT] = "LOAD_ATTR_METHOD_WITH_DICT",
     [EXTENDED_ARG] = "EXTENDED_ARG",
     [LIST_APPEND] = "LIST_APPEND",
     [SET_ADD] = "SET_ADD",
@@ -393,26 +394,27 @@ static const char *const _PyOpcode_OpName[263] = {
     [YIELD_VALUE] = "YIELD_VALUE",
     [RESUME] = "RESUME",
     [MATCH_CLASS] = "MATCH_CLASS",
+    [LOAD_ATTR_METHOD_WITH_VALUES] = "LOAD_ATTR_METHOD_WITH_VALUES",
     [LOAD_CONST__LOAD_FAST] = "LOAD_CONST__LOAD_FAST",
-    [LOAD_FAST__LOAD_CONST] = "LOAD_FAST__LOAD_CONST",
     [FORMAT_VALUE] = "FORMAT_VALUE",
     [BUILD_CONST_KEY_MAP] = "BUILD_CONST_KEY_MAP",
     [BUILD_STRING] = "BUILD_STRING",
+    [LOAD_FAST__LOAD_CONST] = "LOAD_FAST__LOAD_CONST",
     [LOAD_FAST__LOAD_FAST] = "LOAD_FAST__LOAD_FAST",
     [LOAD_GLOBAL_ADAPTIVE] = "LOAD_GLOBAL_ADAPTIVE",
     [LOAD_GLOBAL_BUILTIN] = "LOAD_GLOBAL_BUILTIN",
-    [LOAD_GLOBAL_MODULE] = "LOAD_GLOBAL_MODULE",
     [LIST_EXTEND] = "LIST_EXTEND",
     [SET_UPDATE] = "SET_UPDATE",
     [DICT_MERGE] = "DICT_MERGE",
     [DICT_UPDATE] = "DICT_UPDATE",
+    [LOAD_GLOBAL_MODULE] = "LOAD_GLOBAL_MODULE",
     [STORE_ATTR_ADAPTIVE] = "STORE_ATTR_ADAPTIVE",
     [STORE_ATTR_INSTANCE_VALUE] = "STORE_ATTR_INSTANCE_VALUE",
     [STORE_ATTR_SLOT] = "STORE_ATTR_SLOT",
     [STORE_ATTR_WITH_HINT] = "STORE_ATTR_WITH_HINT",
-    [STORE_FAST__LOAD_FAST] = "STORE_FAST__LOAD_FAST",
     [CALL] = "CALL",
     [KW_NAMES] = "KW_NAMES",
+    [STORE_FAST__LOAD_FAST] = "STORE_FAST__LOAD_FAST",
     [STORE_FAST__STORE_FAST] = "STORE_FAST__STORE_FAST",
     [STORE_SUBSCR_ADAPTIVE] = "STORE_SUBSCR_ADAPTIVE",
     [STORE_SUBSCR_DICT] = "STORE_SUBSCR_DICT",
@@ -421,7 +423,6 @@ static const char *const _PyOpcode_OpName[263] = {
     [UNPACK_SEQUENCE_LIST] = "UNPACK_SEQUENCE_LIST",
     [UNPACK_SEQUENCE_TUPLE] = "UNPACK_SEQUENCE_TUPLE",
     [UNPACK_SEQUENCE_TWO_TUPLE] = "UNPACK_SEQUENCE_TWO_TUPLE",
-    [181] = "<181>",
     [182] = "<182>",
     [183] = "<183>",
     [184] = "<184>",
@@ -507,7 +508,6 @@ static const char *const _PyOpcode_OpName[263] = {
 #endif
 
 #define EXTRA_CASES \
-    case 181: \
     case 182: \
     case 183: \
     case 184: \

Include/opcode.h

@@ -168,37 +168,38 @@ extern "C" {
 #define FOR_ITER_ADAPTIVE                       66
 #define FOR_ITER_LIST                           67
 #define FOR_ITER_RANGE                          72
-#define LOAD_ATTR_ADAPTIVE                      73
+#define FOR_ITER_GEN                            73
-#define LOAD_ATTR_CLASS                         76
+#define LOAD_ATTR_ADAPTIVE                      76
-#define LOAD_ATTR_GETATTRIBUTE_OVERRIDDEN       77
+#define LOAD_ATTR_CLASS                         77
-#define LOAD_ATTR_INSTANCE_VALUE                78
+#define LOAD_ATTR_GETATTRIBUTE_OVERRIDDEN       78
-#define LOAD_ATTR_MODULE                        79
+#define LOAD_ATTR_INSTANCE_VALUE                79
-#define LOAD_ATTR_PROPERTY                      80
+#define LOAD_ATTR_MODULE                        80
-#define LOAD_ATTR_SLOT                          81
+#define LOAD_ATTR_PROPERTY                      81
-#define LOAD_ATTR_WITH_HINT                     86
+#define LOAD_ATTR_SLOT                          86
-#define LOAD_ATTR_METHOD_LAZY_DICT             113
+#define LOAD_ATTR_WITH_HINT                    113
-#define LOAD_ATTR_METHOD_NO_DICT               121
+#define LOAD_ATTR_METHOD_LAZY_DICT             121
-#define LOAD_ATTR_METHOD_WITH_DICT             141
+#define LOAD_ATTR_METHOD_NO_DICT               141
-#define LOAD_ATTR_METHOD_WITH_VALUES           143
+#define LOAD_ATTR_METHOD_WITH_DICT             143
-#define LOAD_CONST__LOAD_FAST                  153
+#define LOAD_ATTR_METHOD_WITH_VALUES           153
-#define LOAD_FAST__LOAD_CONST                  154
+#define LOAD_CONST__LOAD_FAST                  154
-#define LOAD_FAST__LOAD_FAST                   158
+#define LOAD_FAST__LOAD_CONST                  158
-#define LOAD_GLOBAL_ADAPTIVE                   159
+#define LOAD_FAST__LOAD_FAST                   159
-#define LOAD_GLOBAL_BUILTIN                    160
+#define LOAD_GLOBAL_ADAPTIVE                   160
-#define LOAD_GLOBAL_MODULE                     161
+#define LOAD_GLOBAL_BUILTIN                    161
-#define STORE_ATTR_ADAPTIVE                    166
+#define LOAD_GLOBAL_MODULE                     166
-#define STORE_ATTR_INSTANCE_VALUE              167
+#define STORE_ATTR_ADAPTIVE                    167
-#define STORE_ATTR_SLOT                        168
+#define STORE_ATTR_INSTANCE_VALUE              168
-#define STORE_ATTR_WITH_HINT                   169
+#define STORE_ATTR_SLOT                        169
-#define STORE_FAST__LOAD_FAST                  170
+#define STORE_ATTR_WITH_HINT                   170
-#define STORE_FAST__STORE_FAST                 173
+#define STORE_FAST__LOAD_FAST                  173
-#define STORE_SUBSCR_ADAPTIVE                  174
+#define STORE_FAST__STORE_FAST                 174
-#define STORE_SUBSCR_DICT                      175
+#define STORE_SUBSCR_ADAPTIVE                  175
-#define STORE_SUBSCR_LIST_INT                  176
+#define STORE_SUBSCR_DICT                      176
-#define UNPACK_SEQUENCE_ADAPTIVE               177
+#define STORE_SUBSCR_LIST_INT                  177
-#define UNPACK_SEQUENCE_LIST                   178
+#define UNPACK_SEQUENCE_ADAPTIVE               178
-#define UNPACK_SEQUENCE_TUPLE                  179
+#define UNPACK_SEQUENCE_LIST                   179
-#define UNPACK_SEQUENCE_TWO_TUPLE              180
+#define UNPACK_SEQUENCE_TUPLE                  180
+#define UNPACK_SEQUENCE_TWO_TUPLE              181
 #define DO_TRACING                             255
 
 #define HAS_ARG(op) ((((op) >= HAVE_ARGUMENT) && (!IS_PSEUDO_OPCODE(op)))\

Lib/opcode.py

@@ -328,6 +328,7 @@ _specializations = {
         "FOR_ITER_ADAPTIVE",
         "FOR_ITER_LIST",
         "FOR_ITER_RANGE",
+        "FOR_ITER_GEN",
     ],
     "LOAD_ATTR": [
         "LOAD_ATTR_ADAPTIVE",

Lib/test/test_generators.py

@@ -306,6 +306,26 @@ class ExceptionTest(unittest.TestCase):
         self.assertEqual(next(g), "done")
         self.assertEqual(sys.exc_info(), (None, None, None))
 
+    def test_nested_gen_except_loop(self):
+        def gen():
+            for i in range(100):
+                self.assertIsInstance(sys.exception(), TypeError)
+                yield "doing"
+
+        def outer():
+            try:
+                raise TypeError
+            except:
+                for x in gen():
+                    yield x
+
+        try:
+            raise ValueError
+        except Exception:
+            for x in outer():
+                self.assertEqual(x, "doing")
+        self.assertEqual(sys.exception(), None)
+
     def test_except_throw_exception_context(self):
         def gen():
             try:

Misc/NEWS.d/next/Core and Builtins/2022-10-28-09-42-51.gh-issue-96793.ucBfWO.rst

@@ -0,0 +1,3 @@
+Add specialization of :opcode:`FOR_ITER` for generators. Saves multiple
+layers of dispatch and checking to get from the :opcode:`FOR_ITER`
+instruction in the caller to the :opcode:`RESUME` in the generator.

Objects/genobject.c

@@ -209,7 +209,8 @@ gen_send_ex2(PyGenObject *gen, PyObject *arg, PyObject **presult,
 
     frame->previous = tstate->cframe->current_frame;
 
-    gen->gi_exc_state.previous_item = tstate->exc_info;
+    _PyErr_StackItem *prev_exc_info = tstate->exc_info;
+    gen->gi_exc_state.previous_item = prev_exc_info;
     tstate->exc_info = &gen->gi_exc_state;
 
     if (exc) {
@@ -220,12 +221,11 @@ gen_send_ex2(PyGenObject *gen, PyObject *arg, PyObject **presult,
     gen->gi_frame_state = FRAME_EXECUTING;
     EVAL_CALL_STAT_INC(EVAL_CALL_GENERATOR);
     result = _PyEval_EvalFrame(tstate, frame, exc);
+    assert(tstate->exc_info == prev_exc_info);
+    assert(gen->gi_exc_state.previous_item == NULL);
     if (gen->gi_frame_state == FRAME_EXECUTING) {
         gen->gi_frame_state = FRAME_COMPLETED;
     }
-    tstate->exc_info = gen->gi_exc_state.previous_item;
-    gen->gi_exc_state.previous_item = NULL;
-
     assert(tstate->cframe->current_frame == frame->previous);
     /* Don't keep the reference to previous any longer than necessary.  It
      * may keep a chain of frames alive or it could create a reference

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();

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,

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;

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
 };

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);