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

1 files changed, 6 insertions, 3 deletions

diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index 718c8f5..5191b57 100644
--- a/Python/bytecodes.c
+++ b/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, ( -- )) {
-            INSTRUMENTED_JUMP(prev_instr, this_instr+1, PY_MONITORING_EVENT_BRANCH_RIGHT);
+        tier1 op(_MONITOR_END_ASYNC_FOR, ( -- )) {
+            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;