Fix branch monitoring for . Both branches in a pair now have a common source and are included in co_branches

Author: markshannon

Include/internal/pycore_opcode_metadata.h

@@ -936,6 +936,15 @@ def with_extended_args(x):
             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,2,3)])
+
 if check_impl_detail(cpython=True) and ctypes is not None:
     py = ctypes.pythonapi
     freefunc = ctypes.CFUNCTYPE(None,ctypes.c_voidp)

Include/opcode_ids.h

@@ -1683,7 +1683,9 @@ async def foo():
 
 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,25 +1694,26 @@ def check_branches(self, func, tool=TEST_TOOL, recorders=BRANCH_OFFSET_RECORDERS
                 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)
-            func()
-            sys.monitoring.set_local_events(tool, func.__code__, 0)
+            sys.monitoring.set_local_events(tool, test_func.__code__, all_events)
+            run_func()
+            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
+                way, name, src, dest = event
                 if "left" in way:
                     self.assertIn((src, dest), lefts)
                 else:
                     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,26 @@ def foo(n=0):
 
         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
+
+        print(list(foo.__code__.co_branches()))
+        self.check_branches(func, foo)
+
 
 class TestLoadSuperAttr(CheckEvents):
     RECORDERS = CallRecorder, LineRecorder, CRaiseRecorder, CReturnRecorder

Lib/_opcode_metadata.py

@@ -632,6 +632,10 @@ makecode(_PyCompile_CodeUnitMetadata *umd, struct assembler *a, PyObject *const_
     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,11 @@ 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) {
+                    // Adjust the offset so that target is the END_SEND
+                    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;
                 }