GH-96793: Implement PEP 479 in bytecode. (GH-99006)

* Handle converting StopIteration to RuntimeError in bytecode.

* Add custom instruction for converting StopIteration into RuntimeError.

1 files changed, 50 insertions, 12 deletions

diff --git a/Python/compile.c b/Python/compile.c
index bc44bd9..065d1b0 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -1160,6 +1160,9 @@ stack_effect(int opcode, int oparg, int jump)
              * if an exception be raised. */
             return jump ? 1 : 0;
 
+        case STOPITERATION_ERROR:
+            return 0;
+
         case PREP_RERAISE_STAR:
              return -1;
         case RERAISE:
@@ -2545,6 +2548,42 @@ compiler_check_debug_args(struct compiler *c, arguments_ty args)
     return 1;
 }
 
+static inline int
+insert_instruction(basicblock *block, int pos, struct instr *instr) {
+    if (basicblock_next_instr(block) < 0) {
+        return -1;
+    }
+    for (int i = block->b_iused - 1; i > pos; i--) {
+        block->b_instr[i] = block->b_instr[i-1];
+    }
+    block->b_instr[pos] = *instr;
+    return 0;
+}
+
+static int
+wrap_in_stopiteration_handler(struct compiler *c)
+{
+    NEW_JUMP_TARGET_LABEL(c, handler);
+
+    /* Insert SETUP_CLEANUP at start */
+    struct instr setup = {
+        .i_opcode = SETUP_CLEANUP,
+        .i_oparg = handler.id,
+        .i_loc = NO_LOCATION,
+        .i_target = NULL,
+    };
+    if (insert_instruction(c->u->u_cfg_builder.g_entryblock, 0, &setup)) {
+        return 0;
+    }
+
+    ADDOP_LOAD_CONST(c, NO_LOCATION, Py_None);
+    ADDOP(c, NO_LOCATION, RETURN_VALUE);
+    USE_LABEL(c, handler);
+    ADDOP(c, NO_LOCATION, STOPITERATION_ERROR);
+    ADDOP_I(c, NO_LOCATION, RERAISE, 1);
+    return 1;
+}
+
 static int
 compiler_function(struct compiler *c, stmt_ty s, int is_async)
 {
@@ -2625,6 +2664,12 @@ compiler_function(struct compiler *c, stmt_ty s, int is_async)
     for (i = docstring ? 1 : 0; i < asdl_seq_LEN(body); i++) {
         VISIT_IN_SCOPE(c, stmt, (stmt_ty)asdl_seq_GET(body, i));
     }
+    if (c->u->u_ste->ste_coroutine || c->u->u_ste->ste_generator) {
+        if (!wrap_in_stopiteration_handler(c)) {
+            compiler_exit_scope(c);
+            return 0;
+        }
+    }
     co = assemble(c, 1);
     qualname = c->u->u_qualname;
     Py_INCREF(qualname);
@@ -5438,6 +5483,11 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,
     if (type != COMP_GENEXP) {
         ADDOP(c, LOC(e), RETURN_VALUE);
     }
+    if (type == COMP_GENEXP) {
+        if (!wrap_in_stopiteration_handler(c)) {
+            goto error_in_scope;
+        }
+    }
 
     co = assemble(c, 1);
     qualname = c->u->u_qualname;
@@ -8484,18 +8534,6 @@ build_cellfixedoffsets(struct compiler *c)
     return fixed;
 }
 
-static inline int
-insert_instruction(basicblock *block, int pos, struct instr *instr) {
-    if (basicblock_next_instr(block) < 0) {
-        return -1;
-    }
-    for (int i = block->b_iused - 1; i > pos; i--) {
-        block->b_instr[i] = block->b_instr[i-1];
-    }
-    block->b_instr[pos] = *instr;
-    return 0;
-}
-
 static int
 insert_prefix_instructions(struct compiler *c, basicblock *entryblock,
                            int *fixed, int nfreevars, int code_flags)