bpo-43683: Handle generator entry in bytecode

Doc/library/dis.rst

@@ -1247,6 +1247,14 @@ All of the following opcodes use their arguments.
 
    .. versionadded:: 3.10
 
+.. opcode:: GEN_START (kind)
+
+    Pops TOS. If TOS was not ``None``, raises an exception. The ``kind``
+    operand corresponds to the type of generator or coroutine and determines
+    the error message. The legal kinds are 0 for generator, 1 for coroutine,
+    and 2 for async generator.
+
+   .. versionadded:: 3.10
 
 .. opcode:: HAVE_ARGUMENT
 

Lib/importlib/_bootstrap_external.py

@@ -315,6 +315,7 @@ def _write_atomic(path, data, mode=0o666):
 #     Python 3.10a2 3433 (RERAISE restores f_lasti if oparg != 0)
 #     Python 3.10a6 3434 (PEP 634: Structural Pattern Matching)
 #     Python 3.10a7 3435 Use instruction offsets (as opposed to byte offsets).
+#     Python 3.10a7 3436 (Add GEN_START bytecode #43683)
 
 #
 # MAGIC must change whenever the bytecode emitted by the compiler may no

Lib/opcode.py

@@ -173,6 +173,7 @@ def jabs_op(name, op):
 def_op('DELETE_FAST', 126)      # Local variable number
 haslocal.append(126)
 
+def_op('GEN_START', 129)        # Kind of generator/coroutine
 def_op('RAISE_VARARGS', 130)    # Number of raise arguments (1, 2, or 3)
 def_op('CALL_FUNCTION', 131)    # #args
 def_op('MAKE_FUNCTION', 132)    # Flags

Misc/NEWS.d/next/Core and Builtins/2021-04-01-12-30-30.bpo-43683.AjxOx2.rst

@@ -0,0 +1,2 @@
+Add GEN_START opcode. Marks start of generator, including async, or coroutine and handles
+sending values to a newly created generator or coroutine.

Objects/genobject.c

@@ -176,27 +176,12 @@ gen_send_ex2(PyGenObject *gen, PyObject *arg, PyObject **presult,
     }
 
     assert(_PyFrame_IsRunnable(f));
-    if (f->f_lasti == -1) {
-        if (arg && arg != Py_None) {
-            const char *msg = "can't send non-None value to a "
-                              "just-started generator";
-            if (PyCoro_CheckExact(gen)) {
-                msg = NON_INIT_CORO_MSG;
-            }
-            else if (PyAsyncGen_CheckExact(gen)) {
-                msg = "can't send non-None value to a "
-                      "just-started async generator";
-            }
-            PyErr_SetString(PyExc_TypeError, msg);
-            return PYGEN_ERROR;
-        }
-    } else {
-        /* Push arg onto the frame's value stack */
-        result = arg ? arg : Py_None;
-        Py_INCREF(result);
-        gen->gi_frame->f_valuestack[gen->gi_frame->f_stackdepth] = result;
-        gen->gi_frame->f_stackdepth++;
-    }
+    assert(f->f_lasti >= 0 || ((unsigned char *)PyBytes_AS_STRING(f->f_code->co_code))[0] == GEN_START);
+    /* Push arg onto the frame's value stack */
+    result = arg ? arg : Py_None;
+    Py_INCREF(result);
+    gen->gi_frame->f_valuestack[gen->gi_frame->f_stackdepth] = result;
+    gen->gi_frame->f_stackdepth++;
 
     /* Generators always return to their most recent caller, not
      * necessarily their creator. */

Python/ceval.c

@@ -2639,6 +2639,30 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, PyFrameObject *f, int throwflag)
             goto exiting;
         }
 
+        case TARGET(GEN_START): {
+            PyObject *none = POP();
+            Py_DECREF(none);
+            if (none != Py_None) {
+                if (oparg > 2) {
+                    _PyErr_SetString(tstate, PyExc_SystemError,
+                        "Illegal kind for GEN_START");
+                }
+                else {
+                    static const char *gen_kind[3] = {
+                        "generator",
+                        "coroutine",
+                        "async generator"
+                    };
+                    _PyErr_Format(tstate, PyExc_TypeError,
+                        "can't send non-None value to a "
+                                "just-started %s",
+                                gen_kind[oparg]);
+                }
+                goto error;
+            }
+            DISPATCH();
+        }
+
         case TARGET(POP_EXCEPT): {
             PyObject *type, *value, *traceback;
             _PyErr_StackItem *exc_info;

Python/compile.c

@@ -1140,6 +1140,8 @@ stack_effect(int opcode, int oparg, int jump)
             return 1;
         case LIST_TO_TUPLE:
             return 0;
+        case GEN_START:
+            return -1;
         case LIST_EXTEND:
         case SET_UPDATE:
         case DICT_MERGE:
@@ -6169,7 +6171,11 @@ stackdepth(struct compiler *c)
     }
 
     sp = stack;
-    stackdepth_push(&sp, entryblock, 0);
+    if (c->u->u_ste->ste_generator || c->u->u_ste->ste_coroutine) {
+        stackdepth_push(&sp, entryblock, 1);
+    } else {
+        stackdepth_push(&sp, entryblock, 0);
+    }
     while (sp != stack) {
         b = *--sp;
         int depth = b->b_startdepth;
@@ -6648,6 +6654,41 @@ optimize_cfg(struct assembler *a, PyObject *consts);
 static int
 ensure_exits_have_lineno(struct compiler *c);
 
+static int
+insert_generator_prefix(struct compiler *c, basicblock *entryblock) {
+
+    int flags = compute_code_flags(c);
+    if (flags < 0) {
+        return -1;
+    }
+    int kind;
+    if (flags & (CO_GENERATOR | CO_COROUTINE | CO_ASYNC_GENERATOR)) {
+        if (flags & CO_COROUTINE) {
+            kind = 1;
+        }
+        else if (flags & CO_ASYNC_GENERATOR) {
+            kind = 2;
+        }
+        else {
+            kind = 0;
+        }
+    }
+    else {
+        return 0;
+    }
+    if (compiler_next_instr(entryblock) < 0) {
+        return -1;
+    }
+    for (int i = entryblock->b_iused-1; i > 0; i--) {
+        entryblock->b_instr[i] = entryblock->b_instr[i-1];
+    }
+    entryblock->b_instr[0].i_opcode = GEN_START;
+    entryblock->b_instr[0].i_oparg = kind;
+    entryblock->b_instr[0].i_lineno = -1;
+    entryblock->b_instr[0].i_target = NULL;
+    return 0;
+}
+
 static PyCodeObject *
 assemble(struct compiler *c, int addNone)
 {
@@ -6685,6 +6726,10 @@ assemble(struct compiler *c, int addNone)
         entryblock = b;
     }
 
+    if (insert_generator_prefix(c, entryblock)) {
+        goto error;
+    }
+
     /* Set firstlineno if it wasn't explicitly set. */
     if (!c->u->u_firstlineno) {
         if (entryblock && entryblock->b_instr && entryblock->b_instr->i_lineno)