GH-128534: Fix behavior of branch monitoring for async for (GH-130847)

* Both branches in a pair now have a common source and are included in co_branches

Include/internal/pycore_magic_number.h

@@ -270,7 +270,8 @@ Known values:
     Python 3.14a5 3615 (CALL_FUNCTION_EX always take a kwargs argument)
     Python 3.14a5 3616 (Remove BINARY_SUBSCR and family. Make them BINARY_OPs)
     Python 3.14a6 3617 (Branch monitoring for async for loops)
-    Python 3.14a6 3618 (Renumber RESUME opcode from 149 to 128)
+    Python 3.14a6 3618 (Add oparg to END_ASYNC_FOR)
+    Python 3.14a6 3619 (Renumber RESUME opcode from 149 to 128)
 
     Python 3.15 will start with 3650
 
@@ -283,7 +284,7 @@ PC/launcher.c must also be updated.
 
 */
 
-#define PYC_MAGIC_NUMBER 3618
+#define PYC_MAGIC_NUMBER 3619
 /* This is equivalent to converting PYC_MAGIC_NUMBER to 2 bytes
    (little-endian) and then appending b'\r\n'. */
 #define PYC_MAGIC_NUMBER_TOKEN \

Include/internal/pycore_opcode_metadata.h

@@ -2084,7 +2084,7 @@ const struct opcode_metadata _PyOpcode_opcode_metadata[266] = {
     [DELETE_SUBSCR] = { true, INSTR_FMT_IX, HAS_ERROR_FLAG | HAS_ESCAPES_FLAG },
     [DICT_MERGE] = { true, INSTR_FMT_IB, HAS_ARG_FLAG | HAS_ERROR_FLAG | HAS_ESCAPES_FLAG },
     [DICT_UPDATE] = { true, INSTR_FMT_IB, HAS_ARG_FLAG | HAS_ERROR_FLAG | HAS_ESCAPES_FLAG },
-    [END_ASYNC_FOR] = { true, INSTR_FMT_IX, HAS_ERROR_FLAG | HAS_ERROR_NO_POP_FLAG | HAS_ESCAPES_FLAG },
+    [END_ASYNC_FOR] = { true, INSTR_FMT_IB, HAS_ARG_FLAG | HAS_JUMP_FLAG | HAS_ERROR_FLAG | HAS_ERROR_NO_POP_FLAG | HAS_ESCAPES_FLAG },
     [END_FOR] = { true, INSTR_FMT_IX, HAS_ESCAPES_FLAG | HAS_NO_SAVE_IP_FLAG },
     [END_SEND] = { true, INSTR_FMT_IX, HAS_ESCAPES_FLAG | HAS_PURE_FLAG },
     [ENTER_EXECUTOR] = { true, INSTR_FMT_IB, HAS_ARG_FLAG },
@@ -2108,7 +2108,7 @@ const struct opcode_metadata _PyOpcode_opcode_metadata[266] = {
     [INSTRUMENTED_CALL] = { true, INSTR_FMT_IBC00, HAS_ARG_FLAG | HAS_EVAL_BREAK_FLAG | HAS_ERROR_FLAG | HAS_ERROR_NO_POP_FLAG | HAS_ESCAPES_FLAG },
     [INSTRUMENTED_CALL_FUNCTION_EX] = { true, INSTR_FMT_IX, HAS_EVAL_BREAK_FLAG | HAS_ERROR_FLAG | HAS_ERROR_NO_POP_FLAG | HAS_ESCAPES_FLAG },
     [INSTRUMENTED_CALL_KW] = { true, INSTR_FMT_IBC00, HAS_ARG_FLAG | HAS_ERROR_FLAG | HAS_ERROR_NO_POP_FLAG | HAS_ESCAPES_FLAG },
-    [INSTRUMENTED_END_ASYNC_FOR] = { true, INSTR_FMT_IX, HAS_ERROR_FLAG | HAS_ERROR_NO_POP_FLAG | HAS_ESCAPES_FLAG },
+    [INSTRUMENTED_END_ASYNC_FOR] = { true, INSTR_FMT_IB, HAS_ARG_FLAG | HAS_JUMP_FLAG | HAS_ERROR_FLAG | HAS_ERROR_NO_POP_FLAG | HAS_ESCAPES_FLAG },
     [INSTRUMENTED_END_FOR] = { true, INSTR_FMT_IX, HAS_ERROR_FLAG | HAS_ERROR_NO_POP_FLAG | HAS_ESCAPES_FLAG | HAS_NO_SAVE_IP_FLAG },
     [INSTRUMENTED_END_SEND] = { true, INSTR_FMT_IX, HAS_ERROR_FLAG | HAS_ERROR_NO_POP_FLAG | HAS_ESCAPES_FLAG },
     [INSTRUMENTED_FOR_ITER] = { true, INSTR_FMT_IBC, HAS_ARG_FLAG | HAS_JUMP_FLAG | HAS_ERROR_FLAG | HAS_ERROR_NO_POP_FLAG | HAS_ESCAPES_FLAG },

Include/opcode_ids.h

@@ -18,65 +18,65 @@ extern "C" {
 #define CHECK_EXC_MATCH                          5
 #define CLEANUP_THROW                            6
 #define DELETE_SUBSCR                            7
-#define END_ASYNC_FOR                            8
+#define END_FOR                                  8
-#define END_FOR                                  9
+#define END_SEND                                 9
-#define END_SEND                                10
+#define EXIT_INIT_CHECK                         10
-#define EXIT_INIT_CHECK                         11
+#define FORMAT_SIMPLE                           11
-#define FORMAT_SIMPLE                           12
+#define FORMAT_WITH_SPEC                        12
-#define FORMAT_WITH_SPEC                        13
+#define GET_AITER                               13
-#define GET_AITER                               14
+#define GET_ANEXT                               14
-#define GET_ANEXT                               15
+#define GET_ITER                                15
-#define GET_ITER                                16
+#define GET_LEN                                 16
 #define RESERVED                                17
-#define GET_LEN                                 18
+#define GET_YIELD_FROM_ITER                     18
-#define GET_YIELD_FROM_ITER                     19
+#define INTERPRETER_EXIT                        19
-#define INTERPRETER_EXIT                        20
+#define LOAD_BUILD_CLASS                        20
-#define LOAD_BUILD_CLASS                        21
+#define LOAD_LOCALS                             21
-#define LOAD_LOCALS                             22
+#define MAKE_FUNCTION                           22
-#define MAKE_FUNCTION                           23
+#define MATCH_KEYS                              23
-#define MATCH_KEYS                              24
+#define MATCH_MAPPING                           24
-#define MATCH_MAPPING                           25
+#define MATCH_SEQUENCE                          25
-#define MATCH_SEQUENCE                          26
+#define NOP                                     26
-#define NOP                                     27
+#define NOT_TAKEN                               27
-#define NOT_TAKEN                               28
+#define POP_EXCEPT                              28
-#define POP_EXCEPT                              29
+#define POP_ITER                                29
-#define POP_ITER                                30
+#define POP_TOP                                 30
-#define POP_TOP                                 31
+#define PUSH_EXC_INFO                           31
-#define PUSH_EXC_INFO                           32
+#define PUSH_NULL                               32
-#define PUSH_NULL                               33
+#define RETURN_GENERATOR                        33
-#define RETURN_GENERATOR                        34
+#define RETURN_VALUE                            34
-#define RETURN_VALUE                            35
+#define SETUP_ANNOTATIONS                       35
-#define SETUP_ANNOTATIONS                       36
+#define STORE_SLICE                             36
-#define STORE_SLICE                             37
+#define STORE_SUBSCR                            37
-#define STORE_SUBSCR                            38
+#define TO_BOOL                                 38
-#define TO_BOOL                                 39
+#define UNARY_INVERT                            39
-#define UNARY_INVERT                            40
+#define UNARY_NEGATIVE                          40
-#define UNARY_NEGATIVE                          41
+#define UNARY_NOT                               41
-#define UNARY_NOT                               42
+#define WITH_EXCEPT_START                       42
-#define WITH_EXCEPT_START                       43
+#define BINARY_OP                               43
-#define BINARY_OP                               44
+#define BUILD_LIST                              44
-#define BUILD_LIST                              45
+#define BUILD_MAP                               45
-#define BUILD_MAP                               46
+#define BUILD_SET                               46
-#define BUILD_SET                               47
+#define BUILD_SLICE                             47
-#define BUILD_SLICE                             48
+#define BUILD_STRING                            48
-#define BUILD_STRING                            49
+#define BUILD_TUPLE                             49
-#define BUILD_TUPLE                             50
+#define CALL                                    50
-#define CALL                                    51
+#define CALL_INTRINSIC_1                        51
-#define CALL_INTRINSIC_1                        52
+#define CALL_INTRINSIC_2                        52
-#define CALL_INTRINSIC_2                        53
+#define CALL_KW                                 53
-#define CALL_KW                                 54
+#define COMPARE_OP                              54
-#define COMPARE_OP                              55
+#define CONTAINS_OP                             55
-#define CONTAINS_OP                             56
+#define CONVERT_VALUE                           56
-#define CONVERT_VALUE                           57
+#define COPY                                    57
-#define COPY                                    58
+#define COPY_FREE_VARS                          58
-#define COPY_FREE_VARS                          59
+#define DELETE_ATTR                             59
-#define DELETE_ATTR                             60
+#define DELETE_DEREF                            60
-#define DELETE_DEREF                            61
+#define DELETE_FAST                             61
-#define DELETE_FAST                             62
+#define DELETE_GLOBAL                           62
-#define DELETE_GLOBAL                           63
+#define DELETE_NAME                             63
-#define DELETE_NAME                             64
+#define DICT_MERGE                              64
-#define DICT_MERGE                              65
+#define DICT_UPDATE                             65
-#define DICT_UPDATE                             66
+#define END_ASYNC_FOR                           66
 #define EXTENDED_ARG                            67
 #define FOR_ITER                                68
 #define GET_AWAITABLE                           69
@@ -243,7 +243,7 @@ extern "C" {
 #define SETUP_WITH                             264
 #define STORE_FAST_MAYBE_NULL                  265
 
-#define HAVE_ARGUMENT                           43
+#define HAVE_ARGUMENT                           42
 #define MIN_SPECIALIZED_OPCODE                 129
 #define MIN_INSTRUMENTED_OPCODE                234
 

Lib/_opcode_metadata.py

@@ -220,64 +220,64 @@ opmap = {
     'CHECK_EXC_MATCH': 5,
     'CLEANUP_THROW': 6,
     'DELETE_SUBSCR': 7,
-    'END_ASYNC_FOR': 8,
+    'END_FOR': 8,
-    'END_FOR': 9,
+    'END_SEND': 9,
-    'END_SEND': 10,
+    'EXIT_INIT_CHECK': 10,
-    'EXIT_INIT_CHECK': 11,
+    'FORMAT_SIMPLE': 11,
-    'FORMAT_SIMPLE': 12,
+    'FORMAT_WITH_SPEC': 12,
-    'FORMAT_WITH_SPEC': 13,
+    'GET_AITER': 13,
-    'GET_AITER': 14,
+    'GET_ANEXT': 14,
-    'GET_ANEXT': 15,
+    'GET_ITER': 15,
-    'GET_ITER': 16,
+    'GET_LEN': 16,
-    'GET_LEN': 18,
+    'GET_YIELD_FROM_ITER': 18,
-    'GET_YIELD_FROM_ITER': 19,
+    'INTERPRETER_EXIT': 19,
-    'INTERPRETER_EXIT': 20,
+    'LOAD_BUILD_CLASS': 20,
-    'LOAD_BUILD_CLASS': 21,
+    'LOAD_LOCALS': 21,
-    'LOAD_LOCALS': 22,
+    'MAKE_FUNCTION': 22,
-    'MAKE_FUNCTION': 23,
+    'MATCH_KEYS': 23,
-    'MATCH_KEYS': 24,
+    'MATCH_MAPPING': 24,
-    'MATCH_MAPPING': 25,
+    'MATCH_SEQUENCE': 25,
-    'MATCH_SEQUENCE': 26,
+    'NOP': 26,
-    'NOP': 27,
+    'NOT_TAKEN': 27,
-    'NOT_TAKEN': 28,
+    'POP_EXCEPT': 28,
-    'POP_EXCEPT': 29,
+    'POP_ITER': 29,
-    'POP_ITER': 30,
+    'POP_TOP': 30,
-    'POP_TOP': 31,
+    'PUSH_EXC_INFO': 31,
-    'PUSH_EXC_INFO': 32,
+    'PUSH_NULL': 32,
-    'PUSH_NULL': 33,
+    'RETURN_GENERATOR': 33,
-    'RETURN_GENERATOR': 34,
+    'RETURN_VALUE': 34,
-    'RETURN_VALUE': 35,
+    'SETUP_ANNOTATIONS': 35,
-    'SETUP_ANNOTATIONS': 36,
+    'STORE_SLICE': 36,
-    'STORE_SLICE': 37,
+    'STORE_SUBSCR': 37,
-    'STORE_SUBSCR': 38,
+    'TO_BOOL': 38,
-    'TO_BOOL': 39,
+    'UNARY_INVERT': 39,
-    'UNARY_INVERT': 40,
+    'UNARY_NEGATIVE': 40,
-    'UNARY_NEGATIVE': 41,
+    'UNARY_NOT': 41,
-    'UNARY_NOT': 42,
+    'WITH_EXCEPT_START': 42,
-    'WITH_EXCEPT_START': 43,
+    'BINARY_OP': 43,
-    'BINARY_OP': 44,
+    'BUILD_LIST': 44,
-    'BUILD_LIST': 45,
+    'BUILD_MAP': 45,
-    'BUILD_MAP': 46,
+    'BUILD_SET': 46,
-    'BUILD_SET': 47,
+    'BUILD_SLICE': 47,
-    'BUILD_SLICE': 48,
+    'BUILD_STRING': 48,
-    'BUILD_STRING': 49,
+    'BUILD_TUPLE': 49,
-    'BUILD_TUPLE': 50,
+    'CALL': 50,
-    'CALL': 51,
+    'CALL_INTRINSIC_1': 51,
-    'CALL_INTRINSIC_1': 52,
+    'CALL_INTRINSIC_2': 52,
-    'CALL_INTRINSIC_2': 53,
+    'CALL_KW': 53,
-    'CALL_KW': 54,
+    'COMPARE_OP': 54,
-    'COMPARE_OP': 55,
+    'CONTAINS_OP': 55,
-    'CONTAINS_OP': 56,
+    'CONVERT_VALUE': 56,
-    'CONVERT_VALUE': 57,
+    'COPY': 57,
-    'COPY': 58,
+    'COPY_FREE_VARS': 58,
-    'COPY_FREE_VARS': 59,
+    'DELETE_ATTR': 59,
-    'DELETE_ATTR': 60,
+    'DELETE_DEREF': 60,
-    'DELETE_DEREF': 61,
+    'DELETE_FAST': 61,
-    'DELETE_FAST': 62,
+    'DELETE_GLOBAL': 62,
-    'DELETE_GLOBAL': 63,
+    'DELETE_NAME': 63,
-    'DELETE_NAME': 64,
+    'DICT_MERGE': 64,
-    'DICT_MERGE': 65,
+    'DICT_UPDATE': 65,
-    'DICT_UPDATE': 66,
+    'END_ASYNC_FOR': 66,
     'EXTENDED_ARG': 67,
     'FOR_ITER': 68,
     'GET_AWAITABLE': 69,
@@ -360,5 +360,5 @@ opmap = {
     'STORE_FAST_MAYBE_NULL': 265,
 }
 
-HAVE_ARGUMENT = 43
+HAVE_ARGUMENT = 42
 MIN_INSTRUMENTED_OPCODE = 234

Lib/dis.py

@@ -52,6 +52,7 @@ STORE_FAST_LOAD_FAST = opmap['STORE_FAST_LOAD_FAST']
 STORE_FAST_STORE_FAST = opmap['STORE_FAST_STORE_FAST']
 IS_OP = opmap['IS_OP']
 CONTAINS_OP = opmap['CONTAINS_OP']
+END_ASYNC_FOR = opmap['END_ASYNC_FOR']
 
 CACHE = opmap["CACHE"]
 
@@ -605,7 +606,8 @@ class ArgResolver:
                 argval = self.offset_from_jump_arg(op, arg, offset)
                 lbl = self.get_label_for_offset(argval)
                 assert lbl is not None
-                argrepr = f"to L{lbl}"
+                preposition = "from" if deop == END_ASYNC_FOR else "to"
+                argrepr = f"{preposition} L{lbl}"
             elif deop in (LOAD_FAST_LOAD_FAST, STORE_FAST_LOAD_FAST, STORE_FAST_STORE_FAST):
                 arg1 = arg >> 4
                 arg2 = arg & 15
@@ -745,7 +747,8 @@ def _parse_exception_table(code):
 
 def _is_backward_jump(op):
     return opname[op] in ('JUMP_BACKWARD',
-                          'JUMP_BACKWARD_NO_INTERRUPT')
+                          'JUMP_BACKWARD_NO_INTERRUPT',
+                          'END_ASYNC_FOR') # Not really a jump, but it has a "target"
 
 def _get_instructions_bytes(code, linestarts=None, line_offset=0, co_positions=None,
                             original_code=None, arg_resolver=None):

Lib/test/test_code.py

@@ -953,6 +953,15 @@ class CodeLocationTest(unittest.TestCase):
             get_line_branches(with_extended_args),
             [(1,2,8)])
 
+        async def afunc():
+            async for letter in async_iter1:
+                2
+            3
+
+        self.assertEqual(
+            get_line_branches(afunc),
+            [(1,1,3)])
+
 if check_impl_detail(cpython=True) and ctypes is not None:
     py = ctypes.pythonapi
     freefunc = ctypes.CFUNCTYPE(None,ctypes.c_voidp)

Lib/test/test_dis.py

@@ -1256,6 +1256,23 @@ class DisTests(DisTestBase):
         except Exception as e:
             self.assertIsNone(e.__context__)
 
+    def test_async_for_presentation(self):
+
+        async def afunc():
+            async for letter in async_iter1:
+                l2
+            l3
+
+        disassembly =  self.get_disassembly(afunc)
+        for line in disassembly.split("\n"):
+            if "END_ASYNC_FOR" in line:
+                break
+        else:
+            self.fail("No END_ASYNC_FOR in disassembly of async for")
+        self.assertNotIn("to", line)
+        self.assertIn("from", line)
+
+
     @staticmethod
     def code_quicken(f):
         _testinternalcapi = import_helper.import_module("_testinternalcapi")

Lib/test/test_monitoring.py

@@ -1683,7 +1683,9 @@ class TestBranchAndJumpEvents(CheckEvents):
 
 class TestBranchConsistency(MonitoringTestBase, unittest.TestCase):
 
-    def check_branches(self, func, tool=TEST_TOOL, recorders=BRANCH_OFFSET_RECORDERS):
+    def check_branches(self, run_func, test_func=None, tool=TEST_TOOL, recorders=BRANCH_OFFSET_RECORDERS):
+        if test_func is None:
+            test_func = run_func
         try:
             self.assertEqual(sys.monitoring._all_events(), {})
             event_list = []
@@ -1692,16 +1694,17 @@ class TestBranchConsistency(MonitoringTestBase, unittest.TestCase):
                 ev = recorder.event_type
                 sys.monitoring.register_callback(tool, ev, recorder(event_list))
                 all_events |= ev
-            sys.monitoring.set_local_events(tool, func.__code__, all_events)
+            sys.monitoring.set_local_events(tool, test_func.__code__, all_events)
-            func()
+            run_func()
-            sys.monitoring.set_local_events(tool, func.__code__, 0)
+            sys.monitoring.set_local_events(tool, test_func.__code__, 0)
             for recorder in recorders:
                 sys.monitoring.register_callback(tool, recorder.event_type, None)
             lefts = set()
             rights = set()
-            for (src, left, right) in func.__code__.co_branches():
+            for (src, left, right) in test_func.__code__.co_branches():
                 lefts.add((src, left))
                 rights.add((src, right))
+            print(event_list)
             for event in event_list:
                 way, _, src, dest = event
                 if "left" in way:
@@ -1710,7 +1713,7 @@ class TestBranchConsistency(MonitoringTestBase, unittest.TestCase):
                     self.assertIn("right", way)
                     self.assertIn((src, dest), rights)
         finally:
-            sys.monitoring.set_local_events(tool, func.__code__, 0)
+            sys.monitoring.set_local_events(tool, test_func.__code__, 0)
             for recorder in recorders:
                 sys.monitoring.register_callback(tool, recorder.event_type, None)
 
@@ -1762,6 +1765,25 @@ class TestBranchConsistency(MonitoringTestBase, unittest.TestCase):
 
         self.check_branches(foo)
 
+    def test_async_for(self):
+
+        async def gen():
+            yield 2
+            yield 3
+
+        async def foo():
+            async for y in gen():
+                2
+            pass # line 3
+
+        def func():
+            try:
+                foo().send(None)
+            except StopIteration:
+                pass
+
+        self.check_branches(func, foo)
+
 
 class TestLoadSuperAttr(CheckEvents):
     RECORDERS = CallRecorder, LineRecorder, CRaiseRecorder, CReturnRecorder

Misc/NEWS.d/next/Core_and_Builtins/2025-03-04-15-12-32.gh-issue-128534.3A0K3D.rst

@@ -0,0 +1,2 @@
+Ensure that both left and right branches have the same source for ``async for`` loops.
+Add these branches to the ``co_branches()`` iterator.

Programs/test_frozenmain.h

@@ -2,18 +2,18 @@
 unsigned char M_test_frozenmain[] = {
     227,0,0,0,0,0,0,0,0,0,0,0,0,9,0,0,
     0,0,0,0,0,243,184,0,0,0,128,0,90,0,80,0,
-    71,0,112,0,90,0,80,0,71,1,112,1,89,2,33,0,
+    71,0,112,0,90,0,80,0,71,1,112,1,89,2,32,0,
-    80,1,51,1,0,0,0,0,0,0,31,0,89,2,33,0,
+    80,1,50,1,0,0,0,0,0,0,30,0,89,2,32,0,
     80,2,89,0,78,6,0,0,0,0,0,0,0,0,0,0,
-    0,0,0,0,0,0,0,0,51,2,0,0,0,0,0,0,
+    0,0,0,0,0,0,0,0,50,2,0,0,0,0,0,0,
-    31,0,89,1,78,8,0,0,0,0,0,0,0,0,0,0,
+    30,0,89,1,78,8,0,0,0,0,0,0,0,0,0,0,
-    0,0,0,0,0,0,0,0,33,0,51,0,0,0,0,0,
+    0,0,0,0,0,0,0,0,32,0,50,0,0,0,0,0,
-    0,0,80,3,44,26,0,0,0,0,0,0,0,0,0,0,
+    0,0,80,3,43,26,0,0,0,0,0,0,0,0,0,0,
-    112,5,80,4,16,0,68,24,0,0,112,6,89,2,33,0,
+    112,5,80,4,15,0,68,24,0,0,112,6,89,2,32,0,
-    80,5,89,6,12,0,80,6,89,5,89,6,44,26,0,0,
+    80,5,89,6,11,0,80,6,89,5,89,6,43,26,0,0,
-    0,0,0,0,0,0,0,0,12,0,49,4,51,1,0,0,
+    0,0,0,0,0,0,0,0,11,0,48,4,50,1,0,0,
-    0,0,0,0,31,0,73,26,0,0,9,0,30,0,80,0,
+    0,0,0,0,30,0,73,26,0,0,8,0,29,0,80,0,
-    35,0,41,7,78,122,18,70,114,111,122,101,110,32,72,101,
+    34,0,41,7,78,122,18,70,114,111,122,101,110,32,72,101,
     108,108,111,32,87,111,114,108,100,122,8,115,121,115,46,97,
     114,103,118,218,6,99,111,110,102,105,103,41,5,218,12,112,
     114,111,103,114,97,109,95,110,97,109,101,218,10,101,120,101,

Python/assemble.c

@@ -632,6 +632,10 @@ error:
     return co;
 }
 
+
+// The offset (in code units) of the END_SEND from the SEND in the `yield from` sequence.
+#define END_SEND_OFFSET 5
+
 static int
 resolve_jump_offsets(instr_sequence *instrs)
 {
@@ -670,7 +674,12 @@ resolve_jump_offsets(instr_sequence *instrs)
             if (OPCODE_HAS_JUMP(instr->i_opcode)) {
                 instruction *target = &instrs->s_instrs[instr->i_target];
                 instr->i_oparg = target->i_offset;
-                if (instr->i_oparg < offset) {
+                if (instr->i_opcode == END_ASYNC_FOR) {
+                    // sys.monitoring needs to be able to find the matching END_SEND
+                    // but the target is the SEND, so we adjust it here.
+                    instr->i_oparg = offset - instr->i_oparg - END_SEND_OFFSET;
+                }
+                else if (instr->i_oparg < offset) {
                     assert(IS_BACKWARDS_JUMP_OPCODE(instr->i_opcode));
                     instr->i_oparg = offset - instr->i_oparg;
                 }

Python/bytecodes.c

@@ -1341,6 +1341,8 @@ dummy_func(
         }
 
         tier1 op(_END_ASYNC_FOR, (awaitable_st, exc_st -- )) {
+            JUMPBY(0); // Pretend jump as we need source offset for monitoring
+            (void)oparg;
             PyObject *exc = PyStackRef_AsPyObjectBorrow(exc_st);
 
             assert(exc && PyExceptionInstance_Check(exc));
@@ -1356,12 +1358,13 @@ dummy_func(
             }
         }
 
-        tier1 op(_MONITOR_BRANCH_RIGHT, ( -- )) {
+        tier1 op(_MONITOR_END_ASYNC_FOR, ( -- )) {
-            INSTRUMENTED_JUMP(prev_instr, this_instr+1, PY_MONITORING_EVENT_BRANCH_RIGHT);
+            assert((next_instr-oparg)->op.code == END_SEND || (next_instr-oparg)->op.code >= MIN_INSTRUMENTED_OPCODE);
+            INSTRUMENTED_JUMP(next_instr-oparg, this_instr+1, PY_MONITORING_EVENT_BRANCH_RIGHT);
         }
 
         macro(INSTRUMENTED_END_ASYNC_FOR) =
-            _MONITOR_BRANCH_RIGHT +
+            _MONITOR_END_ASYNC_FOR +
             _END_ASYNC_FOR;
 
         macro(END_ASYNC_FOR) = _END_ASYNC_FOR;

Python/codegen.c

@@ -2019,13 +2019,13 @@ codegen_for(compiler *c, stmt_ty s)
     return SUCCESS;
 }
 
-
 static int
 codegen_async_for(compiler *c, stmt_ty s)
 {
     location loc = LOC(s);
 
     NEW_JUMP_TARGET_LABEL(c, start);
+    NEW_JUMP_TARGET_LABEL(c, send);
     NEW_JUMP_TARGET_LABEL(c, except);
     NEW_JUMP_TARGET_LABEL(c, end);
 
@@ -2039,6 +2039,7 @@ codegen_async_for(compiler *c, stmt_ty s)
     ADDOP_JUMP(c, loc, SETUP_FINALLY, except);
     ADDOP(c, loc, GET_ANEXT);
     ADDOP_LOAD_CONST(c, loc, Py_None);
+    USE_LABEL(c, send);
     ADD_YIELD_FROM(c, loc, 1);
     ADDOP(c, loc, POP_BLOCK);  /* for SETUP_FINALLY */
     ADDOP(c, loc, NOT_TAKEN);
@@ -2057,7 +2058,7 @@ codegen_async_for(compiler *c, stmt_ty s)
     /* Use same line number as the iterator,
      * as the END_ASYNC_FOR succeeds the `for`, not the body. */
     loc = LOC(s->v.AsyncFor.iter);
-    ADDOP(c, loc, END_ASYNC_FOR);
+    ADDOP_JUMP(c, loc, END_ASYNC_FOR, send);
 
     /* `else` block */
     VISIT_SEQ(c, stmt, s->v.AsyncFor.orelse);
@@ -4252,6 +4253,7 @@ codegen_async_comprehension_generator(compiler *c, location loc,
                                       int iter_on_stack)
 {
     NEW_JUMP_TARGET_LABEL(c, start);
+    NEW_JUMP_TARGET_LABEL(c, send);
     NEW_JUMP_TARGET_LABEL(c, except);
     NEW_JUMP_TARGET_LABEL(c, if_cleanup);
 
@@ -4279,6 +4281,7 @@ codegen_async_comprehension_generator(compiler *c, location loc,
     ADDOP_JUMP(c, loc, SETUP_FINALLY, except);
     ADDOP(c, loc, GET_ANEXT);
     ADDOP_LOAD_CONST(c, loc, Py_None);
+    USE_LABEL(c, send);
     ADD_YIELD_FROM(c, loc, 1);
     ADDOP(c, loc, POP_BLOCK);
     VISIT(c, expr, gen->target);
@@ -4338,7 +4341,7 @@ codegen_async_comprehension_generator(compiler *c, location loc,
 
     USE_LABEL(c, except);
 
-    ADDOP(c, loc, END_ASYNC_FOR);
+    ADDOP_JUMP(c, loc, END_ASYNC_FOR, send);
 
     return SUCCESS;
 }

Python/flowgraph.c

@@ -849,7 +849,7 @@ calculate_stackdepth(cfg_builder *g)
                 goto error;
             }
             maxdepth = Py_MAX(maxdepth, depth + effects.max);
-            if (HAS_TARGET(instr->i_opcode)) {
+            if (HAS_TARGET(instr->i_opcode) && instr->i_opcode != END_ASYNC_FOR) {
                 if (get_stack_effects(instr->i_opcode, instr->i_oparg, 1, &effects) < 0) {
                     PyErr_Format(PyExc_SystemError,
                                  "Invalid stack effect for opcode=%d, arg=%i",

Python/generated_cases.c.h

@@ -5183,6 +5183,8 @@
             _PyStackRef exc_st;
             exc_st = stack_pointer[-1];
             awaitable_st = stack_pointer[-2];
+            JUMPBY(0); // Pretend jump as we need source offset for monitoring
+            (void)oparg;
             PyObject *exc = PyStackRef_AsPyObjectBorrow(exc_st);
             assert(exc && PyExceptionInstance_Check(exc));
             _PyFrame_SetStackPointer(frame, stack_pointer);
@@ -6607,7 +6609,6 @@
             int opcode = INSTRUMENTED_END_ASYNC_FOR;
             (void)(opcode);
             #endif
-            _Py_CODEUNIT* const prev_instr = frame->instr_ptr;
             _Py_CODEUNIT* const this_instr = next_instr;
             (void)this_instr;
             frame->instr_ptr = next_instr;
@@ -6615,14 +6616,17 @@
             INSTRUCTION_STATS(INSTRUMENTED_END_ASYNC_FOR);
             _PyStackRef awaitable_st;
             _PyStackRef exc_st;
-            // _MONITOR_BRANCH_RIGHT
+            // _MONITOR_END_ASYNC_FOR
             {
-                INSTRUMENTED_JUMP(prev_instr, this_instr+1, PY_MONITORING_EVENT_BRANCH_RIGHT);
+                assert((next_instr-oparg)->op.code == END_SEND || (next_instr-oparg)->op.code >= MIN_INSTRUMENTED_OPCODE);
+                INSTRUMENTED_JUMP(next_instr-oparg, this_instr+1, PY_MONITORING_EVENT_BRANCH_RIGHT);
             }
             // _END_ASYNC_FOR
             {
                 exc_st = stack_pointer[-1];
                 awaitable_st = stack_pointer[-2];
+                JUMPBY(0); // Pretend jump as we need source offset for monitoring
+                (void)oparg;
                 PyObject *exc = PyStackRef_AsPyObjectBorrow(exc_st);
                 assert(exc && PyExceptionInstance_Check(exc));
                 _PyFrame_SetStackPointer(frame, stack_pointer);

Python/instrumentation.c

@@ -3109,6 +3109,14 @@ branchesiter_next(branchesiterator *bi)
                 int not_taken = next_offset + 1;
                 bi->bi_offset = not_taken;
                 return int_triple(offset*2, not_taken*2, (next_offset + oparg)*2);
+            case END_ASYNC_FOR:
+                oparg = (oparg << 8) | inst.op.arg;
+                int src_offset = next_offset - oparg;
+                bi->bi_offset = next_offset;
+                assert(_Py_GetBaseCodeUnit(bi->bi_code, src_offset).op.code == END_SEND);
+                assert(_Py_GetBaseCodeUnit(bi->bi_code, src_offset+1).op.code == NOT_TAKEN);
+                not_taken = src_offset + 2;
+                return int_triple(src_offset *2, not_taken*2, next_offset*2);
             default:
                 oparg = 0;
         }

Python/opcode_targets.h

@@ -8,7 +8,6 @@ static void *opcode_targets[256] = {
     &&TARGET_CHECK_EXC_MATCH,
     &&TARGET_CLEANUP_THROW,
     &&TARGET_DELETE_SUBSCR,
-    &&TARGET_END_ASYNC_FOR,
     &&TARGET_END_FOR,
     &&TARGET_END_SEND,
     &&TARGET_EXIT_INIT_CHECK,
@@ -17,8 +16,8 @@ static void *opcode_targets[256] = {
     &&TARGET_GET_AITER,
     &&TARGET_GET_ANEXT,
     &&TARGET_GET_ITER,
-    &&TARGET_RESERVED,
     &&TARGET_GET_LEN,
+    &&TARGET_RESERVED,
     &&TARGET_GET_YIELD_FROM_ITER,
     &&TARGET_INTERPRETER_EXIT,
     &&TARGET_LOAD_BUILD_CLASS,
@@ -67,6 +66,7 @@ static void *opcode_targets[256] = {
     &&TARGET_DELETE_NAME,
     &&TARGET_DICT_MERGE,
     &&TARGET_DICT_UPDATE,
+    &&TARGET_END_ASYNC_FOR,
     &&TARGET_EXTENDED_ARG,
     &&TARGET_FOR_ITER,
     &&TARGET_GET_AWAITABLE,