bpo-33041: Rework compiling an "async for" loop. (#6142)

* Added new opcode END_ASYNC_FOR.
* Setting global StopAsyncIteration no longer breaks "async for" loops.
* Jumping into an "async for" loop is now disabled.
* Jumping out of an "async for" loop no longer corrupts the stack.
* Simplify the compiler.

Doc/library/dis.rst

@@ -588,6 +588,17 @@ the original TOS1.
    .. versionadded:: 3.5
 
 
+.. opcode:: END_ASYNC_FOR
+
+   Terminates an :keyword:`async for` loop.  Handles an exception raised
+   when awaiting a next item.  If TOS is :exc:`StopAsyncIteration` pop 7
+   values from the stack and restore the exception state using the second
+   three of them.  Otherwise re-raise the exception using the three values
+   from the stack.  An exception handler block is removed from the block stack.
+
+   .. versionadded:: 3.8
+
+
 .. opcode:: BEFORE_ASYNC_WITH
 
    Resolves ``__aenter__`` and ``__aexit__`` from the object on top of the

Doc/whatsnew/3.8.rst

@@ -157,3 +157,7 @@ CPython bytecode changes
 
   (Contributed by Mark Shannon, Antoine Pitrou and Serhiy Storchaka in
   :issue:`17611`.)
+
+* Added new opcode :opcode:`END_ASYNC_FOR` for handling exceptions raised
+  when awaiting a next item in an :keyword:`async for` loop.
+  (Contributed by Serhiy Storchaka in :issue:`33041`.)

Include/opcode.h

@@ -34,6 +34,7 @@ extern "C" {
 #define GET_ANEXT                51
 #define BEFORE_ASYNC_WITH        52
 #define BEGIN_FINALLY            53
+#define END_ASYNC_FOR            54
 #define INPLACE_ADD              55
 #define INPLACE_SUBTRACT         56
 #define INPLACE_MULTIPLY         57

Lib/importlib/_bootstrap_external.py

@@ -247,6 +247,7 @@ _code_type = type(_write_atomic.__code__)
 #     Python 3.7a4  3392 (PEP 552: Deterministic pycs #31650)
 #     Python 3.7b1  3393 (remove STORE_ANNOTATION opcode #32550)
 #     Python 3.8a1  3400 (move frame block handling to compiler #17611)
+#     Python 3.8a1  3401 (add END_ASYNC_FOR #33041)
 #
 # MAGIC must change whenever the bytecode emitted by the compiler may no
 # longer be understood by older implementations of the eval loop (usually
@@ -255,7 +256,7 @@ _code_type = type(_write_atomic.__code__)
 # Whenever MAGIC_NUMBER is changed, the ranges in the magic_values array
 # in PC/launcher.c must also be updated.
 
-MAGIC_NUMBER = (3400).to_bytes(2, 'little') + b'\r\n'
+MAGIC_NUMBER = (3401).to_bytes(2, 'little') + b'\r\n'
 _RAW_MAGIC_NUMBER = int.from_bytes(MAGIC_NUMBER, 'little')  # For import.c
 
 _PYCACHE = '__pycache__'

Lib/opcode.py

@@ -88,7 +88,7 @@ def_op('GET_AITER', 50)
 def_op('GET_ANEXT', 51)
 def_op('BEFORE_ASYNC_WITH', 52)
 def_op('BEGIN_FINALLY', 53)
-
+def_op('END_ASYNC_FOR', 54)
 def_op('INPLACE_ADD', 55)
 def_op('INPLACE_SUBTRACT', 56)
 def_op('INPLACE_MULTIPLY', 57)

Lib/test/test_coroutines.py

@@ -1846,6 +1846,36 @@ class CoroutineTest(unittest.TestCase):
             run_async(run_gen()),
             ([], [121]))
 
+    def test_comp_4_2(self):
+        async def f(it):
+            for i in it:
+                yield i
+
+        async def run_list():
+            return [i + 10 async for i in f(range(5)) if 0 < i < 4]
+        self.assertEqual(
+            run_async(run_list()),
+            ([], [11, 12, 13]))
+
+        async def run_set():
+            return {i + 10 async for i in f(range(5)) if 0 < i < 4}
+        self.assertEqual(
+            run_async(run_set()),
+            ([], {11, 12, 13}))
+
+        async def run_dict():
+            return {i + 10: i + 100 async for i in f(range(5)) if 0 < i < 4}
+        self.assertEqual(
+            run_async(run_dict()),
+            ([], {11: 101, 12: 102, 13: 103}))
+
+        async def run_gen():
+            gen = (i + 10 async for i in f(range(5)) if 0 < i < 4)
+            return [g + 100 async for g in gen]
+        self.assertEqual(
+            run_async(run_gen()),
+            ([], [111, 112, 113]))
+
     def test_comp_5(self):
         async def f(it):
             for i in it:

Lib/test/test_dis.py

@@ -747,8 +747,7 @@ Constants:
    1: 1
 Names:
    0: b
-   1: StopAsyncIteration
-   2: c
+   1: c
 Variable names:
    0: a
    1: d"""

Lib/test/test_sys_settrace.py

@@ -33,6 +33,10 @@ class asynctracecontext:
     async def __aexit__(self, *exc_info):
         self.output.append(-self.value)
 
+async def asynciter(iterable):
+    """Convert an iterable to an asynchronous iterator."""
+    for x in iterable:
+        yield x
 
 
 # A very basic example.  If this fails, we're in deep trouble.
@@ -720,6 +724,23 @@ class JumpTestCase(unittest.TestCase):
             output.append(6)
         output.append(7)
 
+    @async_jump_test(4, 5, [3, 5])
+    async def test_jump_out_of_async_for_block_forwards(output):
+        for i in [1]:
+            async for i in asynciter([1, 2]):
+                output.append(3)
+                output.append(4)
+            output.append(5)
+
+    @async_jump_test(5, 2, [2, 4, 2, 4, 5, 6])
+    async def test_jump_out_of_async_for_block_backwards(output):
+        for i in [1]:
+            output.append(2)
+            async for i in asynciter([1]):
+                output.append(4)
+                output.append(5)
+            output.append(6)
+
     @jump_test(1, 2, [3])
     def test_jump_to_codeless_line(output):
         output.append(1)
@@ -1030,6 +1051,17 @@ class JumpTestCase(unittest.TestCase):
             output.append(7)
         output.append(8)
 
+    @async_jump_test(1, 7, [7, 8])
+    async def test_jump_over_async_for_block_before_else(output):
+        output.append(1)
+        if not output:  # always false
+            async for i in asynciter([3]):
+                output.append(4)
+        else:
+            output.append(6)
+            output.append(7)
+        output.append(8)
+
     # The second set of 'jump' tests are for things that are not allowed:
 
     @jump_test(2, 3, [1], (ValueError, 'after'))
@@ -1081,12 +1113,24 @@ class JumpTestCase(unittest.TestCase):
         for i in 1, 2:
             output.append(3)
 
+    @async_jump_test(1, 3, [], (ValueError, 'into'))
+    async def test_no_jump_forwards_into_async_for_block(output):
+        output.append(1)
+        async for i in asynciter([1, 2]):
+            output.append(3)
+
     @jump_test(3, 2, [2, 2], (ValueError, 'into'))
     def test_no_jump_backwards_into_for_block(output):
         for i in 1, 2:
             output.append(2)
         output.append(3)
 
+    @async_jump_test(3, 2, [2, 2], (ValueError, 'into'))
+    async def test_no_jump_backwards_into_async_for_block(output):
+        async for i in asynciter([1, 2]):
+            output.append(2)
+        output.append(3)
+
     @jump_test(1, 3, [], (ValueError, 'into'))
     def test_no_jump_forwards_into_with_block(output):
         output.append(1)
@@ -1220,6 +1264,17 @@ class JumpTestCase(unittest.TestCase):
             output.append(7)
         output.append(8)
 
+    @async_jump_test(7, 4, [1, 6], (ValueError, 'into'))
+    async def test_no_jump_into_async_for_block_before_else(output):
+        output.append(1)
+        if not output:  # always false
+            async for i in asynciter([3]):
+                output.append(4)
+        else:
+            output.append(6)
+            output.append(7)
+        output.append(8)
+
     def test_no_jump_to_non_integers(self):
         self.run_test(no_jump_to_non_integers, 2, "Spam", [True])
 

Misc/NEWS.d/next/Core and Builtins/2018-03-18-13-56-14.bpo-33041.XwPhI2.rst

+Added new opcode :opcode:`END_ASYNC_FOR` and fixes the following issues:
+
+* Setting global :exc:`StopAsyncIteration` no longer breaks ``async for``
+  loops.
+* Jumping into an ``async for`` loop is now disabled.
+* Jumping out of an ``async for`` loop no longer corrupts the stack.

Objects/frameobject.c

@@ -100,8 +100,7 @@ frame_setlineno(PyFrameObject *f, PyObject* p_new_lineno)
     int line = 0;                       /* (ditto) */
     int addr = 0;                       /* (ditto) */
     int delta_iblock = 0;               /* Scanning the SETUPs and POPs */
-    int for_loop_delta = 0;             /* (ditto) */
-    int delta;
+    int delta = 0;
     int blockstack[CO_MAXBLOCKS];       /* Walking the 'finally' blocks */
     int blockstack_top = 0;             /* (ditto) */
 
@@ -256,14 +255,16 @@ frame_setlineno(PyFrameObject *f, PyObject* p_new_lineno)
                 return -1;
             }
             if (first_in && !second_in) {
-                if (op == FOR_ITER && !delta_iblock) {
-                    for_loop_delta++;
-                }
-                if (op != FOR_ITER) {
+                if (op != FOR_ITER && code[target_addr] != END_ASYNC_FOR) {
                     delta_iblock++;
                 }
+                else if (!delta_iblock) {
+                    /* Pop the iterators of any 'for' and 'async for' loop
+                     * we're jumping out of. */
+                    delta++;
+                }
             }
-            if (op != FOR_ITER) {
+            if (op != FOR_ITER && code[target_addr] != END_ASYNC_FOR) {
                 blockstack[blockstack_top++] = target_addr;
             }
             break;
@@ -289,11 +290,10 @@ frame_setlineno(PyFrameObject *f, PyObject* p_new_lineno)
     assert(blockstack_top == 0);
 
     /* Pop any blocks that we're jumping out of. */
-    delta = 0;
     if (delta_iblock > 0) {
         f->f_iblock -= delta_iblock;
         PyTryBlock *b = &f->f_blockstack[f->f_iblock];
-        delta = (f->f_stacktop - f->f_valuestack) - b->b_level;
+        delta += (f->f_stacktop - f->f_valuestack) - b->b_level;
         if (b->b_type == SETUP_FINALLY &&
             code[b->b_handler] == WITH_CLEANUP_START)
         {
@@ -301,9 +301,6 @@ frame_setlineno(PyFrameObject *f, PyObject* p_new_lineno)
             delta++;
         }
     }
-    /* Pop the iterators of any 'for' loop we're jumping out of. */
-    delta += for_loop_delta;
-
     while (delta > 0) {
         PyObject *v = (*--f->f_stacktop);
         Py_DECREF(v);

Python/ceval.c

@@ -1944,6 +1944,26 @@ main_loop:
             }
         }
 
+        TARGET(END_ASYNC_FOR) {
+            PyObject *exc = POP();
+            assert(PyExceptionClass_Check(exc));
+            if (PyErr_GivenExceptionMatches(exc, PyExc_StopAsyncIteration)) {
+                PyTryBlock *b = PyFrame_BlockPop(f);
+                assert(b->b_type == EXCEPT_HANDLER);
+                Py_DECREF(exc);
+                UNWIND_EXCEPT_HANDLER(b);
+                Py_DECREF(POP());
+                JUMPBY(oparg);
+                FAST_DISPATCH();
+            }
+            else {
+                PyObject *val = POP();
+                PyObject *tb = POP();
+                PyErr_Restore(exc, val, tb);
+                goto exception_unwind;
+            }
+        }
+
         TARGET(LOAD_BUILD_CLASS) {
             _Py_IDENTIFIER(__build_class__);
 

Python/compile.c

@@ -1111,6 +1111,8 @@ stack_effect(int opcode, int oparg, int jump)
             return 1;
         case GET_YIELD_FROM_ITER:
             return 0;
+        case END_ASYNC_FOR:
+            return -7;
         case FORMAT_VALUE:
             /* If there's a fmt_spec on the stack, we go from 2->1,
                else 1->1. */
@@ -2434,78 +2436,40 @@ compiler_for(struct compiler *c, stmt_ty s)
 static int
 compiler_async_for(struct compiler *c, stmt_ty s)
 {
-    _Py_IDENTIFIER(StopAsyncIteration);
-
-    basicblock *try, *except, *end, *after_try, *try_cleanup,
-               *after_loop_else;
-
-    PyObject *stop_aiter_error = _PyUnicode_FromId(&PyId_StopAsyncIteration);
-    if (stop_aiter_error == NULL) {
-        return 0;
-    }
-
-    try = compiler_new_block(c);
+    basicblock *start, *except, *end;
+    start = compiler_new_block(c);
     except = compiler_new_block(c);
     end = compiler_new_block(c);
-    after_try = compiler_new_block(c);
-    try_cleanup = compiler_new_block(c);
-    after_loop_else = compiler_new_block(c);
-
-    if (try == NULL || except == NULL || end == NULL
-            || after_try == NULL || try_cleanup == NULL
-            || after_loop_else == NULL)
-        return 0;
 
-    if (!compiler_push_fblock(c, FOR_LOOP, try, end))
+    if (start == NULL || except == NULL || end == NULL)
         return 0;
 
     VISIT(c, expr, s->v.AsyncFor.iter);
     ADDOP(c, GET_AITER);
 
-    compiler_use_next_block(c, try);
-
+    compiler_use_next_block(c, start);
+    if (!compiler_push_fblock(c, FOR_LOOP, start, end))
+        return 0;
 
     /* SETUP_FINALLY to guard the __anext__ call */
     ADDOP_JREL(c, SETUP_FINALLY, except);
-    if (!compiler_push_fblock(c, EXCEPT, try, NULL))
-        return 0;
-
     ADDOP(c, GET_ANEXT);
     ADDOP_O(c, LOAD_CONST, Py_None, consts);
     ADDOP(c, YIELD_FROM);
     ADDOP(c, POP_BLOCK);  /* for SETUP_FINALLY */
+
+    /* Success block for __anext__ */
     VISIT(c, expr, s->v.AsyncFor.target);
-    compiler_pop_fblock(c, EXCEPT, try);
-    ADDOP_JREL(c, JUMP_FORWARD, after_try);
+    VISIT_SEQ(c, stmt, s->v.AsyncFor.body);
+    ADDOP_JABS(c, JUMP_ABSOLUTE, start);
 
+    compiler_pop_fblock(c, FOR_LOOP, start);
 
     /* Except block for __anext__ */
     compiler_use_next_block(c, except);
-    ADDOP(c, DUP_TOP);
-    ADDOP_O(c, LOAD_GLOBAL, stop_aiter_error, names);
-    ADDOP_I(c, COMPARE_OP, PyCmp_EXC_MATCH);
-    ADDOP_JABS(c, POP_JUMP_IF_FALSE, try_cleanup);
-
-    ADDOP(c, POP_TOP);
-    ADDOP(c, POP_TOP);
-    ADDOP(c, POP_TOP);
-    ADDOP(c, POP_EXCEPT); /* for SETUP_FINALLY */
-    ADDOP(c, POP_TOP);  /* pop iterator from stack */
-    ADDOP_JABS(c, JUMP_ABSOLUTE, after_loop_else);
-
-
-    compiler_use_next_block(c, try_cleanup);
-    ADDOP(c, END_FINALLY);
-
-    /* Success block for __anext__ */
-    compiler_use_next_block(c, after_try);
-    VISIT_SEQ(c, stmt, s->v.AsyncFor.body);
-    ADDOP_JABS(c, JUMP_ABSOLUTE, try);
-
-    compiler_pop_fblock(c, FOR_LOOP, try);
+    ADDOP(c, END_ASYNC_FOR);
 
     /* `else` block */
-    compiler_use_next_block(c, after_loop_else);
     VISIT_SEQ(c, stmt, s->v.For.orelse);
 
     compiler_use_next_block(c, end);
@@ -4003,28 +3967,14 @@ compiler_async_comprehension_generator(struct compiler *c,
                                       asdl_seq *generators, int gen_index,
                                       expr_ty elt, expr_ty val, int type)
 {
-    _Py_IDENTIFIER(StopAsyncIteration);
-
     comprehension_ty gen;
-    basicblock *anchor, *if_cleanup, *try,
-               *after_try, *except, *try_cleanup;
+    basicblock *start, *if_cleanup, *except;
     Py_ssize_t i, n;
-
-    PyObject *stop_aiter_error = _PyUnicode_FromId(&PyId_StopAsyncIteration);
-    if (stop_aiter_error == NULL) {
-        return 0;
-    }
-
-    try = compiler_new_block(c);
-    after_try = compiler_new_block(c);
-    try_cleanup = compiler_new_block(c);
+    start = compiler_new_block(c);
     except = compiler_new_block(c);
     if_cleanup = compiler_new_block(c);
-    anchor = compiler_new_block(c);
 
-    if (if_cleanup == NULL || anchor == NULL ||
-            try == NULL || after_try == NULL ||
-            except == NULL || try_cleanup == NULL) {
+    if (start == NULL || if_cleanup == NULL || except == NULL) {
         return 0;
     }
 
@@ -4041,39 +3991,14 @@ compiler_async_comprehension_generator(struct compiler *c,
         ADDOP(c, GET_AITER);
     }
 
-    compiler_use_next_block(c, try);
-
+    compiler_use_next_block(c, start);
 
     ADDOP_JREL(c, SETUP_FINALLY, except);
-    if (!compiler_push_fblock(c, EXCEPT, try, NULL))
-        return 0;
-
     ADDOP(c, GET_ANEXT);
     ADDOP_O(c, LOAD_CONST, Py_None, consts);
     ADDOP(c, YIELD_FROM);
     ADDOP(c, POP_BLOCK);
     VISIT(c, expr, gen->target);
-    compiler_pop_fblock(c, EXCEPT, try);
-    ADDOP_JREL(c, JUMP_FORWARD, after_try);
-
-
-    compiler_use_next_block(c, except);
-    ADDOP(c, DUP_TOP);
-    ADDOP_O(c, LOAD_GLOBAL, stop_aiter_error, names);
-    ADDOP_I(c, COMPARE_OP, PyCmp_EXC_MATCH);
-    ADDOP_JABS(c, POP_JUMP_IF_FALSE, try_cleanup);
-
-    ADDOP(c, POP_TOP);
-    ADDOP(c, POP_TOP);
-    ADDOP(c, POP_TOP);
-    ADDOP(c, POP_EXCEPT); /* for SETUP_FINALLY */
-    ADDOP_JABS(c, JUMP_ABSOLUTE, anchor);
-
-
-    compiler_use_next_block(c, try_cleanup);
-    ADDOP(c, END_FINALLY);
-
-    compiler_use_next_block(c, after_try);
 
     n = asdl_seq_LEN(gen->ifs);
     for (i = 0; i < n; i++) {
@@ -4118,9 +4043,10 @@ compiler_async_comprehension_generator(struct compiler *c,
         }
     }
     compiler_use_next_block(c, if_cleanup);
-    ADDOP_JABS(c, JUMP_ABSOLUTE, try);
-    compiler_use_next_block(c, anchor);
-    ADDOP(c, POP_TOP);
+    ADDOP_JABS(c, JUMP_ABSOLUTE, start);
+
+    compiler_use_next_block(c, except);
+    ADDOP(c, END_ASYNC_FOR);
 
     return 1;
 }

Python/opcode_targets.h

@@ -53,7 +53,7 @@ static void *opcode_targets[256] = {
     &&TARGET_GET_ANEXT,
     &&TARGET_BEFORE_ASYNC_WITH,
     &&TARGET_BEGIN_FINALLY,
-    &&_unknown_opcode,
+    &&TARGET_END_ASYNC_FOR,
     &&TARGET_INPLACE_ADD,
     &&TARGET_INPLACE_SUBTRACT,
     &&TARGET_INPLACE_MULTIPLY,