diff --git a/Tests/test_async.py b/Tests/test_async.py new file mode 100644 index 000000000..857fc9a13 --- /dev/null +++ b/Tests/test_async.py @@ -0,0 +1,28 @@ +import gc +import pyjion +import unittest +import asyncio + + +class AsyncGeneratorTestCase(unittest.TestCase): + + def setUp(self) -> None: + pyjion.enable() + + def tearDown(self) -> None: + pyjion.disable() + gc.collect() + + def test_basic_generator(self): + async def yield_values(): + for n in range(10): + yield n + + async def calculate_generator(): + total = 1 + async for i in yield_values(): + total += i + return total + + result = asyncio.run(calculate_generator()) + self.assertEqual(result, 46) diff --git a/src/pyjion/absint.cpp b/src/pyjion/absint.cpp index 3906455f4..db7ee1fa0 100644 --- a/src/pyjion/absint.cpp +++ b/src/pyjion/absint.cpp @@ -78,9 +78,8 @@ AbstractInterpreter::~AbstractInterpreter() { } AbstractInterpreterResult AbstractInterpreter::preprocess() { - if (mCode->co_flags & (CO_COROUTINE | CO_ITERABLE_COROUTINE | CO_ASYNC_GENERATOR)) { - // Don't compile co-routines or generators. We can't rely on - // detecting yields because they could be optimized out. + if(mCode->co_flags & CO_ASYNC_GENERATOR){ + // not supported. return IncompatibleCompilerFlags; } @@ -1627,6 +1626,13 @@ void AbstractInterpreter::yieldValue(py_opindex index, size_t stackSize, Instruc for (size_t i = (stackSize - 1); i > 0 ; --i) { m_comp->emit_store_in_frame_value_stack(i-1); } + + if (mCode->co_flags & CO_ASYNC_GENERATOR) { + m_comp->emit_wrap_async_generator(); + errorCheck("failed to wrap async generator"); + m_comp->emit_store_local(m_retValue); + } + m_comp->emit_set_stacktop(stackSize-1); m_comp->emit_branch(BranchAlways, m_retLabel); // ^ Exit Frame || 🔽 Enter frame from next() diff --git a/src/pyjion/dis.py b/src/pyjion/dis.py index 94e230a3c..c4a1dc701 100644 --- a/src/pyjion/dis.py +++ b/src/pyjion/dis.py @@ -141,8 +141,7 @@ class MethodTokens (Enum): METHOD_ISNOT = 0x0000004A METHOD_IS_BOOL = 0x0000004B METHOD_ISNOT_BOOL = 0x0000004C - METHOD_GETITER_OPTIMIZED_TOKEN = 0x0000004D - METHOD_COMPARE_EXCEPTIONS_INT = 0x0000004E + METHOD_WRAP_ASYNC_GENERATOR = 0x0000004D METHOD_UNARY_NOT_INT = 0x00000051 METHOD_FLOAT_FROM_DOUBLE = 0x00000053 diff --git a/src/pyjion/intrins.cpp b/src/pyjion/intrins.cpp index 4f6299b82..973dc0659 100644 --- a/src/pyjion/intrins.cpp +++ b/src/pyjion/intrins.cpp @@ -2651,4 +2651,8 @@ void PyJit_PgcGuardException(PyObject* obj, const char* expected) { expected, PyUnicode_AsUTF8(PyObject_Repr(obj)), obj->ob_type->tp_name); -} \ No newline at end of file +} + +PyObject* PyJit_WrapAsyncGenerator(PyObject *val){ + // TODO: implement +} diff --git a/src/pyjion/intrins.h b/src/pyjion/intrins.h index 0bae9021e..9d2b719ab 100644 --- a/src/pyjion/intrins.h +++ b/src/pyjion/intrins.h @@ -247,6 +247,8 @@ PyObject* PyJit_PyTuple_New(int32_t len); PyObject* PyJit_BuildClass(PyFrameObject *f); +PyObject* PyJit_WrapAsyncGenerator(PyObject *val); + PyObject* PyJit_LoadAttr(PyObject* owner, PyObject* name); PyObject* PyJit_LoadAttrHash(PyObject* owner, PyObject* key, Py_hash_t name_hash); diff --git a/src/pyjion/ipycomp.h b/src/pyjion/ipycomp.h index 80223f26b..000dae716 100644 --- a/src/pyjion/ipycomp.h +++ b/src/pyjion/ipycomp.h @@ -422,6 +422,8 @@ class IPythonCompiler { virtual void emit_debug_msg(const char* msg) = 0; virtual void emit_debug_pyobject() = 0; + virtual void emit_wrap_async_generator() = 0; + // Python 3.7 method calls virtual void emit_load_method(void* name) = 0; diff --git a/src/pyjion/pycomp.cpp b/src/pyjion/pycomp.cpp index fe37e2ac1..c2a74d11c 100644 --- a/src/pyjion/pycomp.cpp +++ b/src/pyjion/pycomp.cpp @@ -2408,6 +2408,9 @@ void PythonCompiler::emit_escape_edges(vector edges, Local success){ } } +void PythonCompiler::emit_wrap_async_generator() { + m_il.emit_call(METHOD_WRAP_ASYNC_GENERATOR); +} /************************************************************************ * End Compiler interface implementation @@ -2647,6 +2650,7 @@ GLOBAL_METHOD(METHOD_SETUP_ANNOTATIONS, &PyJit_SetupAnnotations, CORINFO_TYPE_IN GLOBAL_METHOD(METHOD_LOAD_ASSERTION_ERROR, &PyJit_LoadAssertionError, CORINFO_TYPE_NATIVEINT); GLOBAL_METHOD(METHOD_DEALLOC_OBJECT, &_Py_Dealloc, CORINFO_TYPE_VOID, Parameter(CORINFO_TYPE_NATIVEINT)); +GLOBAL_METHOD(METHOD_WRAP_ASYNC_GENERATOR, &PyJit_WrapAsyncGenerator, CORINFO_TYPE_NATIVEINT, Parameter(CORINFO_TYPE_NATIVEINT)); GLOBAL_METHOD(METHOD_TRACE_LINE, &PyJit_TraceLine, CORINFO_TYPE_VOID, Parameter(CORINFO_TYPE_NATIVEINT), Parameter(CORINFO_TYPE_NATIVEINT), Parameter(CORINFO_TYPE_NATIVEINT), Parameter(CORINFO_TYPE_NATIVEINT)); GLOBAL_METHOD(METHOD_TRACE_FRAME_ENTRY, &PyJit_TraceFrameEntry, CORINFO_TYPE_VOID, Parameter(CORINFO_TYPE_NATIVEINT), ); diff --git a/src/pyjion/pycomp.h b/src/pyjion/pycomp.h index ded218978..04b5fe7fb 100644 --- a/src/pyjion/pycomp.h +++ b/src/pyjion/pycomp.h @@ -117,6 +117,7 @@ #define METHOD_ISNOT 0x0000004A #define METHOD_IS_BOOL 0x0000004B #define METHOD_ISNOT_BOOL 0x0000004C +#define METHOD_WRAP_ASYNC_GENERATOR 0x0000004D #define METHOD_FLOAT_FROM_DOUBLE 0x00000053 #define METHOD_BOOL_FROM_LONG 0x00000054 @@ -504,7 +505,7 @@ class PythonCompiler : public IPythonCompiler { void emit_set_stacktop(size_t height) override; void emit_init_stacktop_local() override; void emit_shrink_stacktop_local(size_t height) override; - + void emit_wrap_async_generator() override; private: void load_frame(); void load_tstate();