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

Author: markshannon

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

@@ -328,6 +328,7 @@ def pseudo_op(name, op, real_ops):
         "FOR_ITER_ADAPTIVE",
         "FOR_ITER_LIST",
         "FOR_ITER_RANGE",
+        "FOR_ITER_GEN",
     ],
     "LOAD_ATTR": [
         "LOAD_ATTR_ADAPTIVE",

Include/opcode.h

@@ -306,6 +306,26 @@ def gen():
         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:

Lib/opcode.py

@@ -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.

Lib/test/test_generators.py

@@ -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

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

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