1 files changed, 31 insertions, 16 deletions

diff --git a/Python/compile.c b/Python/compile.c
index 5915a4bc2f8..b4e06e7cb14 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -7527,6 +7527,9 @@ build_cellfixedoffsets(_PyCompile_CodeUnitMetadata *umd)
     return fixed;
 }
 
+#define IS_GENERATOR(CF) \
+    ((CF) & (CO_GENERATOR | CO_COROUTINE | CO_ASYNC_GENERATOR))
+
 static int
 insert_prefix_instructions(_PyCompile_CodeUnitMetadata *umd, basicblock *entryblock,
                            int *fixed, int nfreevars, int code_flags)
@@ -7534,7 +7537,11 @@ insert_prefix_instructions(_PyCompile_CodeUnitMetadata *umd, basicblock *entrybl
     assert(umd->u_firstlineno > 0);
 
     /* Add the generator prefix instructions. */
-    if (code_flags & (CO_GENERATOR | CO_COROUTINE | CO_ASYNC_GENERATOR)) {
+    if (IS_GENERATOR(code_flags)) {
+        /* Note that RETURN_GENERATOR + POP_TOP have a net stack effect
+         * of 0. This is because RETURN_GENERATOR pushes an element
+         * with _PyFrame_StackPush before switching stacks.
+         */
         cfg_instr make_gen = {
             .i_opcode = RETURN_GENERATOR,
             .i_oparg = 0,
@@ -7715,19 +7722,27 @@ optimize_and_assemble_code_unit(struct compiler_unit *u, PyObject *const_cache,
     int nparams = (int)PyList_GET_SIZE(u->u_ste->ste_varnames);
     int nlocals = (int)PyDict_GET_SIZE(u->u_metadata.u_varnames);
     assert(u->u_metadata.u_firstlineno);
-    if (_PyCfg_OptimizeCodeUnit(&g, consts, const_cache, code_flags, nlocals,
+    if (_PyCfg_OptimizeCodeUnit(&g, consts, const_cache, nlocals,
                                 nparams, u->u_metadata.u_firstlineno) < 0) {
         goto error;
     }
 
-    /** Assembly **/
-    int nlocalsplus = prepare_localsplus(&u->u_metadata, &g, code_flags);
-    if (nlocalsplus < 0) {
+    int stackdepth = _PyCfg_Stackdepth(&g);
+    if (stackdepth < 0) {
         goto error;
     }
 
-    int maxdepth = _PyCfg_Stackdepth(g.g_entryblock, code_flags);
-    if (maxdepth < 0) {
+    /* prepare_localsplus adds instructions for generators that push
+     * and pop an item on the stack. This assertion makes sure there
+     * is space on the stack for that.
+     * It should always be true, because at least one expression is
+     * required to turn a function into a generator.
+     */
+    assert(!(IS_GENERATOR(code_flags) && stackdepth == 0));
+
+    /** Assembly **/
+    int nlocalsplus = prepare_localsplus(&u->u_metadata, &g, code_flags);
+    if (nlocalsplus < 0) {
         goto error;
     }
 
@@ -7746,7 +7761,7 @@ optimize_and_assemble_code_unit(struct compiler_unit *u, PyObject *const_cache,
     }
 
     co = _PyAssemble_MakeCodeObject(&u->u_metadata, const_cache, consts,
-                                    maxdepth, &optimized_instrs, nlocalsplus,
+                                    stackdepth, &optimized_instrs, nlocalsplus,
                                     code_flags, filename);
 
 error:
@@ -8190,8 +8205,8 @@ _PyCompile_OptimizeCfg(PyObject *instructions, PyObject *consts, int nlocals)
     if (instructions_to_cfg(instructions, &g) < 0) {
         goto error;
     }
-    int code_flags = 0, nparams = 0, firstlineno = 1;
-    if (_PyCfg_OptimizeCodeUnit(&g, consts, const_cache, code_flags, nlocals,
+    int nparams = 0, firstlineno = 1;
+    if (_PyCfg_OptimizeCodeUnit(&g, consts, const_cache, nlocals,
                                 nparams, firstlineno) < 0) {
         goto error;
     }
@@ -8226,14 +8241,14 @@ _PyCompile_Assemble(_PyCompile_CodeUnitMetadata *umd, PyObject *filename,
         goto error;
     }
 
-    int code_flags = 0;
-    int nlocalsplus = prepare_localsplus(umd, &g, code_flags);
-    if (nlocalsplus < 0) {
+    int stackdepth = _PyCfg_Stackdepth(&g);
+    if (stackdepth < 0) {
         goto error;
     }
 
-    int maxdepth = _PyCfg_Stackdepth(g.g_entryblock, code_flags);
-    if (maxdepth < 0) {
+    int code_flags = 0;
+    int nlocalsplus = prepare_localsplus(umd, &g, code_flags);
+    if (nlocalsplus < 0) {
         goto error;
     }
 
@@ -8256,7 +8271,7 @@ _PyCompile_Assemble(_PyCompile_CodeUnitMetadata *umd, PyObject *filename,
         goto error;
     }
     co = _PyAssemble_MakeCodeObject(umd, const_cache,
-                                    consts, maxdepth, &optimized_instrs,
+                                    consts, stackdepth, &optimized_instrs,
                                     nlocalsplus, code_flags, filename);
     Py_DECREF(consts);