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

* Handle converting StopIteration to RuntimeError in bytecode.

* Add custom instruction for converting StopIteration into RuntimeError.

4 files changed, 147 insertions, 24 deletions

diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index 94a4e76..b0d5627 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -1085,6 +1085,49 @@ dummy_func(
             goto exception_unwind;
         }
 
+        inst(STOPITERATION_ERROR) {
+            assert(frame->owner == FRAME_OWNED_BY_GENERATOR);
+            PyObject *exc = TOP();
+            assert(PyExceptionInstance_Check(exc));
+            const char *msg = NULL;
+            if (PyErr_GivenExceptionMatches(exc, PyExc_StopIteration)) {
+                msg = "generator raised StopIteration";
+                if (frame->f_code->co_flags & CO_ASYNC_GENERATOR) {
+                    msg = "async generator raised StopIteration";
+                }
+                else if (frame->f_code->co_flags & CO_COROUTINE) {
+                    msg = "coroutine raised StopIteration";
+                }
+            }
+            else if ((frame->f_code->co_flags & CO_ASYNC_GENERATOR) &&
+                    PyErr_GivenExceptionMatches(exc, PyExc_StopAsyncIteration))
+            {
+                /* code in `gen` raised a StopAsyncIteration error:
+                raise a RuntimeError.
+                */
+                msg = "async generator raised StopAsyncIteration";
+            }
+            if (msg != NULL) {
+                PyObject *message = _PyUnicode_FromASCII(msg, strlen(msg));
+                if (message == NULL) {
+                    goto error;
+                }
+                PyObject *error = PyObject_CallOneArg(PyExc_RuntimeError, message);
+                if (error == NULL) {
+                    Py_DECREF(message);
+                    goto error;
+                }
+                assert(PyExceptionInstance_Check(error));
+                SET_TOP(error);
+                PyException_SetCause(error, exc);
+                Py_INCREF(exc);
+                PyException_SetContext(error, exc);
+                Py_DECREF(message);
+            }
+            DISPATCH();
+        }
+
+
         // stack effect: ( -- __0)
         inst(LOAD_ASSERTION_ERROR) {
             PyObject *value = PyExc_AssertionError;
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)
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
index 36c7498..bea51d7 100644
--- a/Python/generated_cases.c.h
+++ b/Python/generated_cases.c.h
@@ -1011,6 +1011,48 @@
             goto exception_unwind;
         }
 
+        TARGET(STOPITERATION_ERROR) {
+            assert(frame->owner == FRAME_OWNED_BY_GENERATOR);
+            PyObject *exc = TOP();
+            assert(PyExceptionInstance_Check(exc));
+            const char *msg = NULL;
+            if (PyErr_GivenExceptionMatches(exc, PyExc_StopIteration)) {
+                msg = "generator raised StopIteration";
+                if (frame->f_code->co_flags & CO_ASYNC_GENERATOR) {
+                    msg = "async generator raised StopIteration";
+                }
+                else if (frame->f_code->co_flags & CO_COROUTINE) {
+                    msg = "coroutine raised StopIteration";
+                }
+            }
+            else if ((frame->f_code->co_flags & CO_ASYNC_GENERATOR) &&
+                    PyErr_GivenExceptionMatches(exc, PyExc_StopAsyncIteration))
+            {
+                /* code in `gen` raised a StopAsyncIteration error:
+                raise a RuntimeError.
+                */
+                msg = "async generator raised StopAsyncIteration";
+            }
+            if (msg != NULL) {
+                PyObject *message = _PyUnicode_FromASCII(msg, strlen(msg));
+                if (message == NULL) {
+                    goto error;
+                }
+                PyObject *error = PyObject_CallOneArg(PyExc_RuntimeError, message);
+                if (error == NULL) {
+                    Py_DECREF(message);
+                    goto error;
+                }
+                assert(PyExceptionInstance_Check(error));
+                SET_TOP(error);
+                PyException_SetCause(error, exc);
+                Py_INCREF(exc);
+                PyException_SetContext(error, exc);
+                Py_DECREF(message);
+            }
+            DISPATCH();
+        }
+
         TARGET(LOAD_ASSERTION_ERROR) {
             PyObject *value = PyExc_AssertionError;
             Py_INCREF(value);
diff --git a/Python/opcode_targets.h b/Python/opcode_targets.h
index 9004bba..a963a7a 100644
--- a/Python/opcode_targets.h
+++ b/Python/opcode_targets.h
@@ -62,30 +62,30 @@ static void *opcode_targets[256] = {
     &&TARGET_STORE_SUBSCR,
     &&TARGET_DELETE_SUBSCR,
     &&TARGET_COMPARE_OP_INT_JUMP,
+    &&TARGET_STOPITERATION_ERROR,
     &&TARGET_COMPARE_OP_STR_JUMP,
     &&TARGET_EXTENDED_ARG_QUICK,
     &&TARGET_FOR_ITER_ADAPTIVE,
     &&TARGET_FOR_ITER_LIST,
-    &&TARGET_FOR_ITER_RANGE,
     &&TARGET_GET_ITER,
     &&TARGET_GET_YIELD_FROM_ITER,
     &&TARGET_PRINT_EXPR,
     &&TARGET_LOAD_BUILD_CLASS,
+    &&TARGET_FOR_ITER_RANGE,
     &&TARGET_LOAD_ATTR_ADAPTIVE,
-    &&TARGET_LOAD_ATTR_CLASS,
     &&TARGET_LOAD_ASSERTION_ERROR,
     &&TARGET_RETURN_GENERATOR,
+    &&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_LOAD_ATTR_WITH_HINT,
     &&TARGET_LIST_TO_TUPLE,
     &&TARGET_RETURN_VALUE,
     &&TARGET_IMPORT_STAR,
     &&TARGET_SETUP_ANNOTATIONS,
-    &&TARGET_LOAD_ATTR_METHOD_LAZY_DICT,
+    &&TARGET_LOAD_ATTR_WITH_HINT,
     &&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_NO_DICT,
+    &&TARGET_LOAD_ATTR_METHOD_LAZY_DICT,
     &&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_WITH_DICT,
+    &&TARGET_LOAD_ATTR_METHOD_NO_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_VALUES,
+    &&TARGET_LOAD_ATTR_METHOD_WITH_DICT,
     &&TARGET_CALL_FUNCTION_EX,
-    &&TARGET_LOAD_CONST__LOAD_FAST,
+    &&TARGET_LOAD_ATTR_METHOD_WITH_VALUES,
     &&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_CONST__LOAD_FAST,
     &&TARGET_LOAD_FAST__LOAD_CONST,
-    &&TARGET_LOAD_FAST__LOAD_FAST,
     &&TARGET_FORMAT_VALUE,
     &&TARGET_BUILD_CONST_KEY_MAP,
     &&TARGET_BUILD_STRING,
+    &&TARGET_LOAD_FAST__LOAD_FAST,
     &&TARGET_LOAD_GLOBAL_ADAPTIVE,
     &&TARGET_LOAD_GLOBAL_BUILTIN,
     &&TARGET_LOAD_GLOBAL_MODULE,
-    &&TARGET_STORE_ATTR_ADAPTIVE,
     &&TARGET_LIST_EXTEND,
     &&TARGET_SET_UPDATE,
     &&TARGET_DICT_MERGE,
     &&TARGET_DICT_UPDATE,
+    &&TARGET_STORE_ATTR_ADAPTIVE,
     &&TARGET_STORE_ATTR_INSTANCE_VALUE,
     &&TARGET_STORE_ATTR_SLOT,
     &&TARGET_STORE_ATTR_WITH_HINT,
     &&TARGET_STORE_FAST__LOAD_FAST,
-    &&TARGET_STORE_FAST__STORE_FAST,
     &&TARGET_CALL,
     &&TARGET_KW_NAMES,
+    &&TARGET_STORE_FAST__STORE_FAST,
     &&TARGET_STORE_SUBSCR_ADAPTIVE,
     &&TARGET_STORE_SUBSCR_DICT,
     &&TARGET_STORE_SUBSCR_LIST_INT,
@@ -253,6 +254,5 @@ static void *opcode_targets[256] = {
     &&_unknown_opcode,
     &&_unknown_opcode,
     &&_unknown_opcode,
-    &&_unknown_opcode,
     &&TARGET_DO_TRACING
 };