Update generator logic

Author: brijkapadia

Include/cpython/funcobject.h

@@ -43,7 +43,7 @@ typedef struct {
     PyObject *func_annotations; /* Annotations, a dict or NULL */
     PyObject *func_annotate;    /* Callable to fill the annotations dictionary */
     PyObject *func_typeparams;  /* Tuple of active type variables or NULL */
-    PyObject *func_old_codes;
+    PyObject *func_old_codes;   /* NULL or list of past code objects */
     vectorcallfunc vectorcall;
     /* Version number for use by specializer.
      * Can set to non-zero when we want to specialize.

Include/internal/pycore_interpframe.h

@@ -160,6 +160,13 @@ static inline void _PyFrame_Copy(_PyInterpreterFrame *src, _PyInterpreterFrame *
     }
 }
 
+/* Generator frames need a strong reference to the code object */
+static inline void _PyFrame_CopyForGenerators(_PyInterpreterFrame *old_frame, _PyInterpreterFrame *gen_frame) {
+    _PyFrame_Copy(old_frame, gen_frame);
+    gen_frame->owner = FRAME_OWNED_BY_GENERATOR;
+    gen_frame->f_executable = PyStackRef_MakeHeapSafe(gen_frame->f_executable);
+}
+
 #ifdef Py_GIL_DISABLED
 static inline void
 _PyFrame_InitializeTLBC(PyThreadState *tstate, _PyInterpreterFrame *frame,

Include/internal/pycore_interpframe_structs.h

@@ -27,7 +27,7 @@ enum _frameowner {
 };
 
 struct _PyInterpreterFrame {
-    _PyStackRef f_executable; /* Deferred or strong reference (code object or None) */
+    _PyStackRef f_executable; /* Borrowed reference (code object or None) */
     struct _PyInterpreterFrame *previous;
     _PyStackRef f_funcobj; /* Deferred or strong reference. Only valid if not on C stack */
     PyObject *f_globals; /* Borrowed reference. Only valid if not on C stack */

Modules/_testinternalcapi/test_cases.c.h

@@ -93,6 +93,11 @@ gen_traverse(PyObject *self, visitproc visit, void *arg)
             return err;
         }
     }
+    else {
+        // We still need to visit the code object when the frame is cleared to
+        // ensure that it's kept alive if the reference is deferred.
+        _Py_VISIT_STACKREF(gen->gi_iframe.f_executable);
+    }
     /* No need to visit cr_origin, because it's just tuples/str/int, so can't
        participate in a reference cycle. */
     Py_VISIT(gen->gi_exc_state.exc_value);
@@ -1171,11 +1176,10 @@ gen_new_with_qualname(PyTypeObject *type, PyFrameObject *f,
     assert(f->f_frame->frame_obj == NULL);
     assert(f->f_frame->owner == FRAME_OWNED_BY_FRAME_OBJECT);
     _PyInterpreterFrame *frame = &gen->gi_iframe;
-    _PyFrame_Copy((_PyInterpreterFrame *)f->_f_frame_data, frame);
+    _PyFrame_CopyForGenerators((_PyInterpreterFrame *)f->_f_frame_data, frame);
     gen->gi_frame_state = FRAME_CREATED;
     assert(frame->frame_obj == f);
     f->f_frame = frame;
-    frame->owner = FRAME_OWNED_BY_GENERATOR;
     assert(PyObject_GC_IsTracked((PyObject *)f));
     Py_DECREF(f);
     gen->gi_weakreflist = NULL;

Objects/genobject.c

@@ -5854,10 +5854,9 @@ dummy_func(
             SAVE_STACK();
             _PyInterpreterFrame *gen_frame = &gen->gi_iframe;
             frame->instr_ptr++;
-            _PyFrame_Copy(frame, gen_frame);
+            _PyFrame_CopyForGenerators(frame, gen_frame);
             assert(frame->frame_obj == NULL);
             gen->gi_frame_state = FRAME_CREATED;
-            gen_frame->owner = FRAME_OWNED_BY_GENERATOR;
             _Py_LeaveRecursiveCallPy(tstate);
             _PyInterpreterFrame *prev = frame->previous;
             _PyThreadState_UpdateLastProfiledFrame(tstate, frame, prev);