Stackless issue #188: enable soft switching for sub-iterators, (async) generators and coroutines

Enable soft switching for iterators/coroutines called by "yield from".
Improve the finalisation of a soft switched YIELD_FROM instruction.
Enable stackless calls of the following methods:
- generator.send() (generator.__next__() was already stackless);
- coroutine.send();
- coroutine_wrapper.__next__() and coroutine_wrapper.send();
- async_generator_asend.__next__() and async_generator_asend.send().

Author: Anselm Kruis

Objects/genobject.c

@@ -159,13 +159,30 @@ gen_dealloc(PyGenObject *gen)
  * was not faster, but considerably slower than this solution.
  */
 
-static PyObject* gen_iternext_callback(PyFrameObject *f, int exc, PyObject *result);
+/* Like gen_send_ex, additionally pass ob3 into the callback function. */
+static PyObject *
+gen_send_ex2(PyGenObject *gen, PyObject *arg, int exc, int closing, PyObject* ob3);
+
+/* callback for (async) generators and coroutines. */
+static PyObject *
+gen_iternext_callback(PyFrameObject *f, int exc, PyObject *result);
+
+/* Additional callback-code for async generators. */
+static PyObject *
+async_gen_asend_send_end(PyAsyncGenObject *gen, PyObject *asend, PyObject *result);
 #endif
 
+
 static PyObject *
 gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing)
 {
 #ifdef STACKLESS
+    return gen_send_ex2(gen, arg, exc, closing, NULL);
+}
+
+static PyObject *
+gen_send_ex2(PyGenObject *gen, PyObject *arg, int exc, int closing, PyObject * ob3)
+{
     STACKLESS_GETARG();
     PyFrameObject *stopframe;
 #endif
@@ -249,6 +266,7 @@ gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing)
     assert(f->f_back->f_back == NULL);
     assert(((PyCFrameObject *)f->f_back)->ob1 == NULL);
     assert(((PyCFrameObject *)f->f_back)->ob2 == NULL);
+    assert(((PyCFrameObject *)f->f_back)->ob3 == NULL);
 
     if (f->f_lasti != -1)
 #else
@@ -295,8 +313,10 @@ gen_send_ex(PyGenObject *gen, PyObject *arg, int exc, int closing)
 
     Py_INCREF(gen);
     Py_XINCREF(arg);
+    Py_XINCREF(ob3);
     ((PyCFrameObject *) f->f_back)->ob1 = (PyObject *) gen;
     ((PyCFrameObject *) f->f_back)->ob2 = arg;
+    ((PyCFrameObject *) f->f_back)->ob3 = ob3;
 
     if (exc)
         result = NULL;
@@ -326,12 +346,14 @@ gen_iternext_callback(PyFrameObject *f, int exc, PyObject *result)
     PyCFrameObject *cf = (PyCFrameObject *) f;
     PyGenObject *gen = (PyGenObject *) cf->ob1;
     PyObject *arg = cf->ob2;
+    PyObject *ob3 = cf->ob3;
 
     /* We hold references to things in the cframe, if we release it
        before we clear the references, they get incorrectly and
        prematurely freed. */
     cf->ob1 = NULL;
     cf->ob2 = NULL;
+    cf->ob3 = NULL;
 
     f = gen->gi_frame;
     /* Check, that this cframe belongs to gen */
@@ -400,8 +422,13 @@ gen_iternext_callback(PyFrameObject *f, int exc, PyObject *result)
         Py_DECREF(f);
     }
 #ifdef STACKLESS
+    if (ob3 && PyAsyncGen_CheckExact(gen)) {
+        result = async_gen_asend_send_end((PyAsyncGenObject *)gen, ob3, result);
+    }
+
     Py_DECREF(gen);
     Py_XDECREF(arg);
+    Py_XDECREF(ob3);
 #endif
     return result;
 }
@@ -839,7 +866,7 @@ static PyMemberDef gen_memberlist[] = {
 };
 
 static PyMethodDef gen_methods[] = {
-    {"send",(PyCFunction)_PyGen_Send, METH_O, send_doc},
+    {"send",(PyCFunction)_PyGen_Send, METH_O | METH_STACKLESS, send_doc},
     {"throw",(PyCFunction)gen_throw, METH_VARARGS, throw_doc},
     {"close",(PyCFunction)gen_close, METH_NOARGS, close_doc},
     {NULL, NULL}        /* Sentinel */
@@ -1092,7 +1119,7 @@ PyDoc_STRVAR(coro_close_doc,
 "close() -> raise GeneratorExit inside coroutine.");
 
 static PyMethodDef coro_methods[] = {
-    {"send",(PyCFunction)_PyGen_Send, METH_O, coro_send_doc},
+    {"send",(PyCFunction)_PyGen_Send, METH_O | METH_STACKLESS, coro_send_doc},
     {"throw",(PyCFunction)gen_throw, METH_VARARGS, coro_throw_doc},
     {"close",(PyCFunction)gen_close, METH_NOARGS, coro_close_doc},
     {NULL, NULL}        /* Sentinel */
@@ -1197,7 +1224,7 @@ coro_wrapper_traverse(PyCoroWrapper *cw, visitproc visit, void *arg)
 }
 
 static PyMethodDef coro_wrapper_methods[] = {
-    {"send",(PyCFunction)coro_wrapper_send, METH_O, coro_send_doc},
+    {"send",(PyCFunction)coro_wrapper_send, METH_O | METH_STACKLESS, coro_send_doc},
     {"throw",(PyCFunction)coro_wrapper_throw, METH_VARARGS, coro_throw_doc},
     {"close",(PyCFunction)coro_wrapper_close, METH_NOARGS, coro_close_doc},
     {NULL, NULL}        /* Sentinel */
@@ -1245,6 +1272,8 @@ PyTypeObject _PyCoroWrapper_Type = {
     0,                                          /* tp_free */
 };
 
+STACKLESS_DECLARE_METHOD(&_PyCoroWrapper_Type, tp_iternext);
+
 static PyObject *
 compute_cr_origin(int origin_depth)
 {
@@ -1666,6 +1695,7 @@ async_gen_asend_traverse(PyAsyncGenASend *o, visitproc visit, void *arg)
 static PyObject *
 async_gen_asend_send(PyAsyncGenASend *o, PyObject *arg)
 {
+    STACKLESS_GETARG();
     PyObject *result;
 
     if (o->ags_state == AWAITABLE_STATE_CLOSED) {
@@ -1680,6 +1710,16 @@ async_gen_asend_send(PyAsyncGenASend *o, PyObject *arg)
         o->ags_state = AWAITABLE_STATE_ITER;
     }
 
+#ifdef STACKLESS
+    if (stackless) {
+        STACKLESS_PROMOTE_ALL();
+        result = gen_send_ex2((PyGenObject*)o->ags_gen, arg, 0, 0, (PyObject *)o);
+        STACKLESS_ASSERT();
+        if (STACKLESS_UNWINDING(result))
+            return result;
+    }
+    else
+#endif
     result = gen_send_ex((PyGenObject*)o->ags_gen, arg, 0, 0);
     result = async_gen_unwrap_value(o->ags_gen, result);
 
@@ -1690,6 +1730,22 @@ async_gen_asend_send(PyAsyncGenASend *o, PyObject *arg)
     return result;
 }
 
+#ifdef STACKLESS
+static PyObject *
+async_gen_asend_send_end(PyAsyncGenObject *gen, PyObject *asend, PyObject *result)
+{
+    PyAsyncGenASend *o = (PyAsyncGenASend *) asend;
+    assert(o->ags_gen == gen);
+
+    result = async_gen_unwrap_value(o->ags_gen, result);
+
+    if (result == NULL) {
+        o->ags_state = AWAITABLE_STATE_CLOSED;
+    }
+
+    return result;
+}
+#endif
 
 static PyObject *
 async_gen_asend_iternext(PyAsyncGenASend *o)
@@ -1728,7 +1784,7 @@ async_gen_asend_close(PyAsyncGenASend *o, PyObject *args)
 
 
 static PyMethodDef async_gen_asend_methods[] = {
-    {"send", (PyCFunction)async_gen_asend_send, METH_O, send_doc},
+    {"send", (PyCFunction)async_gen_asend_send, METH_O | METH_STACKLESS, send_doc},
     {"throw", (PyCFunction)async_gen_asend_throw, METH_VARARGS, throw_doc},
     {"close", (PyCFunction)async_gen_asend_close, METH_NOARGS, close_doc},
     {NULL, NULL}        /* Sentinel */
@@ -1784,6 +1840,8 @@ PyTypeObject _PyAsyncGenASend_Type = {
     0,                                          /* tp_new */
 };
 
+STACKLESS_DECLARE_METHOD(&_PyAsyncGenASend_Type, tp_iternext);
+
 
 static PyObject *
 async_gen_asend_new(PyAsyncGenObject *gen, PyObject *sendval)

Python/ceval.c

@@ -1100,7 +1100,35 @@ _PyEval_EvalFrameDefault(PyFrameObject *f, int throwflag)
                 /*f->f_execute = PyEval_EvalFrame_value;
                 PREDICT(WITH_CLEANUP_FINISH); */
             }
-        } else
+        }
+        else if (f->f_execute == slp_eval_frame_yield_from) {
+            /* finalise the YIELD_FROM operation */
+            PyObject *receiver = TOP();
+            if (retval == NULL) {
+                int err;
+                if (tstate->c_tracefunc != NULL
+                    && PyErr_ExceptionMatches(PyExc_StopIteration))
+                    call_exc_trace(tstate->c_tracefunc, tstate->c_traceobj, tstate, f);
+                err = _PyGen_FetchStopIterationValue(&retval);
+                if (err < 0) {
+                    retval = NULL;
+                }
+                else {
+                    Py_DECREF(receiver);
+                    SET_TOP(retval);
+                }
+            }
+            else {
+                /* receiver remains on stack, retval is value to be yielded */
+                f->f_stacktop = stack_pointer;
+                why = WHY_YIELD;
+                /* and repeat... */
+                assert(f->f_lasti >= (int)sizeof(_Py_CODEUNIT));
+                f->f_lasti -= sizeof(_Py_CODEUNIT);
+                goto fast_yield;
+            }
+        }
+        else
             Py_FatalError("invalid frame function");
         f->f_execute = slp_eval_frame_value;
     }
@@ -1972,15 +2000,29 @@ _PyEval_EvalFrameDefault(PyFrameObject *f, int throwflag)
             PyObject *receiver = TOP();
             int err;
             if (PyGen_CheckExact(receiver) || PyCoro_CheckExact(receiver)) {
+                STACKLESS_PROPOSE_ALL(tstate);
                 retval = _PyGen_Send((PyGenObject *)receiver, v);
+                STACKLESS_ASSERT();
             } else {
                 _Py_IDENTIFIER(send);
-                if (v == Py_None)
+                if (v == Py_None) {
+                    STACKLESS_PROPOSE_METHOD(tstate, receiver, tp_iternext);
                     retval = Py_TYPE(receiver)->tp_iternext(receiver);
-                else
+                    STACKLESS_ASSERT();
+                }
+                else {
+                    STACKLESS_PROPOSE_ALL(tstate);
                     retval = _PyObject_CallMethodIdObjArgs(receiver, &PyId_send, v, NULL);
+                    STACKLESS_ASSERT();
+                }
             }
             Py_DECREF(v);
+#ifdef STACKLESS
+            if (STACKLESS_UNWINDING(retval)) {
+                HANDLE_UNWINDING(slp_eval_frame_yield_from, 0, retval);
+                receiver = TOP();
+            }
+#endif
             if (retval == NULL) {
                 PyObject *val;
                 if (tstate->c_tracefunc != NULL
@@ -3831,6 +3873,21 @@ slp_eval_frame_with_cleanup(PyFrameObject *f, int throwflag, PyObject *retval)
     return r;
 }
 
+PyObject *
+slp_eval_frame_yield_from(PyFrameObject *f, int throwflag, PyObject *retval)
+{
+    PyObject *r;
+    /*
+     * this function is identical to PyEval_EvalFrame_value.
+     * it serves as a marker whether we are inside of a
+     * WITH_CLEANUP operation.
+     * NOTE / XXX: see above.
+     */
+    SLP_DO_NOT_OPTIMIZE_AWAY((char *)5);
+    r = slp_eval_frame_value(f, throwflag, retval);
+    return r;
+}
+
 static PyObject *
 run_frame_dispatch(PyFrameObject *f, int exc, PyObject *retval)
 {
@@ -3948,7 +4005,8 @@ _PyEval_EvalFrameDefault(PyFrameObject *f, int throwflag)
                 f->f_execute != slp_eval_frame_noval &&
                 f->f_execute != slp_eval_frame_iter &&
                 f->f_execute != slp_eval_frame_setup_with &&
-                f->f_execute != slp_eval_frame_with_cleanup)  {
+                f->f_execute != slp_eval_frame_with_cleanup &&
+                f->f_execute != slp_eval_frame_yield_from)  {
             PyErr_BadInternalCall();
             return NULL;
         }

Stackless/changelog.txt

@@ -9,6 +9,17 @@ What's New in Stackless 3.X.X?
 
 *Release date: 20XX-XX-XX*
 
+- https://github.com/stackless-dev/stackless/issues/188
+  Enable stackless calls of sub-iterators and coroutines (technically: the
+  YIELD_FROM opcode) and of the following methods, if soft-switching is
+  enabled:
+    - generator.send() (generator.__next__() was already stackless);
+    - coroutine.send();
+    - coroutine_wrapper.__next__() and coroutine_wrapper.send();
+    - async_generator_asend.__next__() and async_generator_asend.send().
+  With this change it is possible to soft switch into/from the compound
+  statements "async for" and "async with".
+
 - https://github.com/stackless-dev/stackless/issues/200
   Fix a bug in the C-API functions PyTasklet_Run_nr() and
   PyTasklet_Switch_nr(). Under exotic conditions the functions could

Stackless/core/stackless_impl.h

@@ -378,6 +378,7 @@ PyObject * slp_eval_frame_noval(struct _frame *f,  int throwflag, PyObject *retv
 PyObject * slp_eval_frame_iter(struct _frame *f,  int throwflag, PyObject *retval);
 PyObject * slp_eval_frame_setup_with(struct _frame *f,  int throwflag, PyObject *retval);
 PyObject * slp_eval_frame_with_cleanup(struct _frame *f,  int throwflag, PyObject *retval);
+PyObject * slp_eval_frame_yield_from(struct _frame *f,  int throwflag, PyObject *retval);
 /* other eval_frame functions from module/scheduling.c */
 PyObject * slp_restore_tracing(PyFrameObject *f, int exc, PyObject *retval);
 /* other eval_frame functions from Objects/typeobject.c */

Stackless/core/stackless_methods.h

@@ -14,8 +14,8 @@ typedef struct {
 #define MFLAG_OFS_IND(meth) MFLAG_OFS(meth) + MFLAG_IND
 
 static _stackless_method _stackless_methtable[] = {
-    /* from methodobject.c */
-    {&PyCFunction_Type,            MFLAG_OFS(tp_call)},
+    /* from classobject.c */
+    {&PyMethod_Type,            MFLAG_OFS(tp_call)},
     /* from descrobject.c */
     {&PyMethodDescr_Type,        MFLAG_OFS(tp_call)},
     {&PyClassMethodDescr_Type,    MFLAG_OFS(tp_call)},
@@ -24,10 +24,12 @@ static _stackless_method _stackless_methtable[] = {
     {&PyFunction_Type,            MFLAG_OFS(tp_call)},
     /* from genobject.c */
     {&PyGen_Type,                MFLAG_OFS(tp_iternext)},
+    {&_PyCoroWrapper_Type,        MFLAG_OFS(tp_iternext)},
+    {&_PyAsyncGenASend_Type,    MFLAG_OFS(tp_iternext)},
+    /* from methodobject.c */
+    {&PyCFunction_Type,            MFLAG_OFS(tp_call)},
     /* from typeobject.c */
     {&PyType_Type,                MFLAG_OFS(tp_call)},
-    /* from classobject.c */
-    {&PyMethod_Type,            MFLAG_OFS(tp_call)},
     /* from channelobject.c */
     {&PyChannel_Type,            MFLAG_OFS(tp_iternext)},
     {0, 0} /* sentinel */

Stackless/pickling/prickelpit.c

@@ -839,6 +839,7 @@ SLP_DEF_INVALID_EXEC(eval_frame_noval)
 SLP_DEF_INVALID_EXEC(eval_frame_iter)
 SLP_DEF_INVALID_EXEC(eval_frame_setup_with)
 SLP_DEF_INVALID_EXEC(eval_frame_with_cleanup)
+SLP_DEF_INVALID_EXEC(eval_frame_yield_from)
 SLP_DEF_INVALID_EXEC(slp_channel_seq_callback)
 SLP_DEF_INVALID_EXEC(slp_restore_tracing)
 SLP_DEF_INVALID_EXEC(slp_tp_init_callback)
@@ -1202,6 +1203,8 @@ static int init_frametype(PyObject * mod)
                              slp_eval_frame_setup_with, SLP_REF_INVALID_EXEC(eval_frame_setup_with))
         || slp_register_execute(&PyFrame_Type, "eval_frame_with_cleanup",
                              slp_eval_frame_with_cleanup, SLP_REF_INVALID_EXEC(eval_frame_with_cleanup))
+        || slp_register_execute(&PyFrame_Type, "eval_frame_yield_from",
+                             slp_eval_frame_yield_from, SLP_REF_INVALID_EXEC(eval_frame_yield_from))
         || slp_register_execute(&PyCFrame_Type, "channel_seq_callback",
                              slp_channel_seq_callback, SLP_REF_INVALID_EXEC(slp_channel_seq_callback))
         || slp_register_execute(&PyCFrame_Type, "slp_restore_tracing",