GH-128534: Instrument branches for async for loops. (GH-130569)

Include/internal/pycore_magic_number.h

@@ -269,6 +269,7 @@ Known values:
     Python 3.14a5 3614 (Add BINARY_OP_EXTEND)
     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.15 will start with 3650
 
@@ -281,7 +282,7 @@ PC/launcher.c must also be updated.
 
 */
 
-#define PYC_MAGIC_NUMBER 3616
+#define PYC_MAGIC_NUMBER 3617
 /* 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

@@ -227,6 +227,8 @@ int _PyOpcode_num_popped(int opcode, int oparg)  {
             return 4;
         case INSTRUMENTED_CALL_KW:
             return 3 + oparg;
+        case INSTRUMENTED_END_ASYNC_FOR:
+            return 2;
         case INSTRUMENTED_END_FOR:
             return 2;
         case INSTRUMENTED_END_SEND:
@@ -700,6 +702,8 @@ int _PyOpcode_num_pushed(int opcode, int oparg)  {
             return 1;
         case INSTRUMENTED_CALL_KW:
             return 1;
+        case INSTRUMENTED_END_ASYNC_FOR:
+            return 0;
         case INSTRUMENTED_END_FOR:
             return 1;
         case INSTRUMENTED_END_SEND:
@@ -1383,6 +1387,10 @@ int _PyOpcode_max_stack_effect(int opcode, int oparg, int *effect)  {
             *effect = Py_MAX(0, -2 - oparg);
             return 0;
         }
+        case INSTRUMENTED_END_ASYNC_FOR: {
+            *effect = 0;
+            return 0;
+        }
         case INSTRUMENTED_END_FOR: {
             *effect = -1;
             return 0;
@@ -2097,6 +2105,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_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 },
@@ -2525,6 +2534,7 @@ const char *_PyOpcode_OpName[266] = {
     [INSTRUMENTED_CALL] = "INSTRUMENTED_CALL",
     [INSTRUMENTED_CALL_FUNCTION_EX] = "INSTRUMENTED_CALL_FUNCTION_EX",
     [INSTRUMENTED_CALL_KW] = "INSTRUMENTED_CALL_KW",
+    [INSTRUMENTED_END_ASYNC_FOR] = "INSTRUMENTED_END_ASYNC_FOR",
     [INSTRUMENTED_END_FOR] = "INSTRUMENTED_END_FOR",
     [INSTRUMENTED_END_SEND] = "INSTRUMENTED_END_SEND",
     [INSTRUMENTED_FOR_ITER] = "INSTRUMENTED_FOR_ITER",
@@ -2787,6 +2797,7 @@ const uint8_t _PyOpcode_Deopt[256] = {
     [INSTRUMENTED_CALL] = INSTRUMENTED_CALL,
     [INSTRUMENTED_CALL_FUNCTION_EX] = INSTRUMENTED_CALL_FUNCTION_EX,
     [INSTRUMENTED_CALL_KW] = INSTRUMENTED_CALL_KW,
+    [INSTRUMENTED_END_ASYNC_FOR] = INSTRUMENTED_END_ASYNC_FOR,
     [INSTRUMENTED_END_FOR] = INSTRUMENTED_END_FOR,
     [INSTRUMENTED_END_SEND] = INSTRUMENTED_END_SEND,
     [INSTRUMENTED_FOR_ITER] = INSTRUMENTED_FOR_ITER,
@@ -2951,7 +2962,6 @@ const uint8_t _PyOpcode_Deopt[256] = {
     case 148: \
     case 232: \
     case 233: \
-    case 234: \
         ;
 struct pseudo_targets {
     uint8_t as_sequence;

Include/opcode_ids.h

@@ -210,20 +210,21 @@ extern "C" {
 #define UNPACK_SEQUENCE_LIST                   229
 #define UNPACK_SEQUENCE_TUPLE                  230
 #define UNPACK_SEQUENCE_TWO_TUPLE              231
-#define INSTRUMENTED_END_FOR                   235
+#define INSTRUMENTED_END_FOR                   234
-#define INSTRUMENTED_POP_ITER                  236
+#define INSTRUMENTED_POP_ITER                  235
-#define INSTRUMENTED_END_SEND                  237
+#define INSTRUMENTED_END_SEND                  236
-#define INSTRUMENTED_FOR_ITER                  238
+#define INSTRUMENTED_FOR_ITER                  237
-#define INSTRUMENTED_INSTRUCTION               239
+#define INSTRUMENTED_INSTRUCTION               238
-#define INSTRUMENTED_JUMP_FORWARD              240
+#define INSTRUMENTED_JUMP_FORWARD              239
-#define INSTRUMENTED_NOT_TAKEN                 241
+#define INSTRUMENTED_NOT_TAKEN                 240
-#define INSTRUMENTED_POP_JUMP_IF_TRUE          242
+#define INSTRUMENTED_POP_JUMP_IF_TRUE          241
-#define INSTRUMENTED_POP_JUMP_IF_FALSE         243
+#define INSTRUMENTED_POP_JUMP_IF_FALSE         242
-#define INSTRUMENTED_POP_JUMP_IF_NONE          244
+#define INSTRUMENTED_POP_JUMP_IF_NONE          243
-#define INSTRUMENTED_POP_JUMP_IF_NOT_NONE      245
+#define INSTRUMENTED_POP_JUMP_IF_NOT_NONE      244
-#define INSTRUMENTED_RESUME                    246
+#define INSTRUMENTED_RESUME                    245
-#define INSTRUMENTED_RETURN_VALUE              247
+#define INSTRUMENTED_RETURN_VALUE              246
-#define INSTRUMENTED_YIELD_VALUE               248
+#define INSTRUMENTED_YIELD_VALUE               247
+#define INSTRUMENTED_END_ASYNC_FOR             248
 #define INSTRUMENTED_LOAD_SUPER_ATTR           249
 #define INSTRUMENTED_CALL                      250
 #define INSTRUMENTED_CALL_KW                   251
@@ -244,7 +245,7 @@ extern "C" {
 
 #define HAVE_ARGUMENT                           43
 #define MIN_SPECIALIZED_OPCODE                 150
-#define MIN_INSTRUMENTED_OPCODE                235
+#define MIN_INSTRUMENTED_OPCODE                234
 
 #ifdef __cplusplus
 }

Lib/_opcode_metadata.py

@@ -328,20 +328,21 @@ opmap = {
     'UNPACK_EX': 114,
     'UNPACK_SEQUENCE': 115,
     'YIELD_VALUE': 116,
-    'INSTRUMENTED_END_FOR': 235,
+    'INSTRUMENTED_END_FOR': 234,
-    'INSTRUMENTED_POP_ITER': 236,
+    'INSTRUMENTED_POP_ITER': 235,
-    'INSTRUMENTED_END_SEND': 237,
+    'INSTRUMENTED_END_SEND': 236,
-    'INSTRUMENTED_FOR_ITER': 238,
+    'INSTRUMENTED_FOR_ITER': 237,
-    'INSTRUMENTED_INSTRUCTION': 239,
+    'INSTRUMENTED_INSTRUCTION': 238,
-    'INSTRUMENTED_JUMP_FORWARD': 240,
+    'INSTRUMENTED_JUMP_FORWARD': 239,
-    'INSTRUMENTED_NOT_TAKEN': 241,
+    'INSTRUMENTED_NOT_TAKEN': 240,
-    'INSTRUMENTED_POP_JUMP_IF_TRUE': 242,
+    'INSTRUMENTED_POP_JUMP_IF_TRUE': 241,
-    'INSTRUMENTED_POP_JUMP_IF_FALSE': 243,
+    'INSTRUMENTED_POP_JUMP_IF_FALSE': 242,
-    'INSTRUMENTED_POP_JUMP_IF_NONE': 244,
+    'INSTRUMENTED_POP_JUMP_IF_NONE': 243,
-    'INSTRUMENTED_POP_JUMP_IF_NOT_NONE': 245,
+    'INSTRUMENTED_POP_JUMP_IF_NOT_NONE': 244,
-    'INSTRUMENTED_RESUME': 246,
+    'INSTRUMENTED_RESUME': 245,
-    'INSTRUMENTED_RETURN_VALUE': 247,
+    'INSTRUMENTED_RETURN_VALUE': 246,
-    'INSTRUMENTED_YIELD_VALUE': 248,
+    'INSTRUMENTED_YIELD_VALUE': 247,
+    'INSTRUMENTED_END_ASYNC_FOR': 248,
     'INSTRUMENTED_LOAD_SUPER_ATTR': 249,
     'INSTRUMENTED_CALL': 250,
     'INSTRUMENTED_CALL_KW': 251,
@@ -360,4 +361,4 @@ opmap = {
 }
 
 HAVE_ARGUMENT = 43
-MIN_INSTRUMENTED_OPCODE = 235
+MIN_INSTRUMENTED_OPCODE = 234

Lib/test/test_monitoring.py

@@ -1657,6 +1657,29 @@ class TestBranchAndJumpEvents(CheckEvents):
             in_loop,
             exit_loop])
 
+    def test_async_for(self):
+
+        def func():
+            async def gen():
+                yield 2
+                yield 3
+
+            async def foo():
+                async for y in gen():
+                    2
+                pass # line 3
+
+            try:
+                foo().send(None)
+            except StopIteration:
+                pass
+
+        self.check_events(func, recorders = BRANCHES_RECORDERS, expected = [
+            ('branch left', 'foo', 1, 1),
+            ('branch left', 'foo', 1, 1),
+            ('branch right', 'foo', 1, 3),
+            ('branch left', 'func', 12, 12)])
+
 
 class TestBranchConsistency(MonitoringTestBase, unittest.TestCase):
 

Misc/NEWS.d/next/Core_and_Builtins/2025-02-26-10-32-48.gh-issue-128534.JEiQex.rst

@@ -0,0 +1,2 @@
+Add branch monitoring (``BRANCH_LEFT`` and ``BRANCH_RIGHT`` events) for
+``async for`` loops.

Python/bytecodes.c

@@ -1339,7 +1339,7 @@ dummy_func(
             goto exception_unwind;
         }
 
-        tier1 inst(END_ASYNC_FOR, (awaitable_st, exc_st -- )) {
+        tier1 op(_END_ASYNC_FOR, (awaitable_st, exc_st -- )) {
             PyObject *exc = PyStackRef_AsPyObjectBorrow(exc_st);
 
             assert(exc && PyExceptionInstance_Check(exc));
@@ -1355,6 +1355,16 @@ dummy_func(
             }
         }
 
+        tier1 op(_MONITOR_BRANCH_RIGHT, ( -- )) {
+            INSTRUMENTED_JUMP(prev_instr, this_instr+1, PY_MONITORING_EVENT_BRANCH_RIGHT);
+        }
+
+        macro(INSTRUMENTED_END_ASYNC_FOR) =
+            _MONITOR_BRANCH_RIGHT +
+            _END_ASYNC_FOR;
+
+        macro(END_ASYNC_FOR) = _END_ASYNC_FOR;
+
         tier1 inst(CLEANUP_THROW, (sub_iter, last_sent_val, exc_value_st -- none, value)) {
             PyObject *exc_value = PyStackRef_AsPyObjectBorrow(exc_value_st);
             #if !Py_TAIL_CALL_INTERP

Python/codegen.c

@@ -2041,6 +2041,7 @@ codegen_async_for(compiler *c, stmt_ty s)
     ADDOP_LOAD_CONST(c, loc, Py_None);
     ADD_YIELD_FROM(c, loc, 1);
     ADDOP(c, loc, POP_BLOCK);  /* for SETUP_FINALLY */
+    ADDOP(c, loc, NOT_TAKEN);
 
     /* Success block for __anext__ */
     VISIT(c, expr, s->v.AsyncFor.target);

Python/generated_cases.c.h

@@ -6600,6 +6600,57 @@
             DISPATCH();
         }
 
+        TARGET(INSTRUMENTED_END_ASYNC_FOR) {
+            #if Py_TAIL_CALL_INTERP
+            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;
+            next_instr += 1;
+            INSTRUCTION_STATS(INSTRUMENTED_END_ASYNC_FOR);
+            _PyStackRef awaitable_st;
+            _PyStackRef exc_st;
+            // _MONITOR_BRANCH_RIGHT
+            {
+                INSTRUMENTED_JUMP(prev_instr, this_instr+1, PY_MONITORING_EVENT_BRANCH_RIGHT);
+            }
+            // _END_ASYNC_FOR
+            {
+                exc_st = stack_pointer[-1];
+                awaitable_st = stack_pointer[-2];
+                PyObject *exc = PyStackRef_AsPyObjectBorrow(exc_st);
+                assert(exc && PyExceptionInstance_Check(exc));
+                _PyFrame_SetStackPointer(frame, stack_pointer);
+                int matches = PyErr_GivenExceptionMatches(exc, PyExc_StopAsyncIteration);
+                stack_pointer = _PyFrame_GetStackPointer(frame);
+                if (matches) {
+                    _PyFrame_SetStackPointer(frame, stack_pointer);
+                    _PyStackRef tmp = exc_st;
+                    exc_st = PyStackRef_NULL;
+                    stack_pointer[-1] = exc_st;
+                    PyStackRef_CLOSE(tmp);
+                    tmp = awaitable_st;
+                    awaitable_st = PyStackRef_NULL;
+                    stack_pointer[-2] = awaitable_st;
+                    PyStackRef_CLOSE(tmp);
+                    stack_pointer = _PyFrame_GetStackPointer(frame);
+                    stack_pointer += -2;
+                    assert(WITHIN_STACK_BOUNDS());
+                }
+                else {
+                    Py_INCREF(exc);
+                    _PyFrame_SetStackPointer(frame, stack_pointer);
+                    _PyErr_SetRaisedException(tstate, exc);
+                    monitor_reraise(tstate, frame, this_instr);
+                    JUMP_TO_LABEL(exception_unwind);
+                }
+            }
+            DISPATCH();
+        }
+
         TARGET(INSTRUMENTED_END_FOR) {
             #if Py_TAIL_CALL_INTERP
             int opcode = INSTRUMENTED_END_FOR;

Python/instrumentation.c

@@ -106,6 +106,7 @@ static const int8_t EVENT_FOR_OPCODE[256] = {
     [INSTRUMENTED_END_SEND] = PY_MONITORING_EVENT_STOP_ITERATION,
     [NOT_TAKEN] = PY_MONITORING_EVENT_BRANCH_LEFT,
     [INSTRUMENTED_NOT_TAKEN] = PY_MONITORING_EVENT_BRANCH_LEFT,
+    [END_ASYNC_FOR] = PY_MONITORING_EVENT_BRANCH_RIGHT,
 };
 
 static const uint8_t DE_INSTRUMENT[256] = {
@@ -127,6 +128,7 @@ static const uint8_t DE_INSTRUMENT[256] = {
     [INSTRUMENTED_END_SEND] = END_SEND,
     [INSTRUMENTED_LOAD_SUPER_ATTR] = LOAD_SUPER_ATTR,
     [INSTRUMENTED_NOT_TAKEN] = NOT_TAKEN,
+    [INSTRUMENTED_END_ASYNC_FOR] = END_ASYNC_FOR,
 };
 
 static const uint8_t INSTRUMENTED_OPCODES[256] = {
@@ -166,6 +168,8 @@ static const uint8_t INSTRUMENTED_OPCODES[256] = {
     [INSTRUMENTED_LOAD_SUPER_ATTR] = INSTRUMENTED_LOAD_SUPER_ATTR,
     [NOT_TAKEN] = INSTRUMENTED_NOT_TAKEN,
     [INSTRUMENTED_NOT_TAKEN] = INSTRUMENTED_NOT_TAKEN,
+    [END_ASYNC_FOR] = INSTRUMENTED_END_ASYNC_FOR,
+    [INSTRUMENTED_END_ASYNC_FOR] = INSTRUMENTED_END_ASYNC_FOR,
 
     [INSTRUMENTED_LINE] = INSTRUMENTED_LINE,
     [INSTRUMENTED_INSTRUCTION] = INSTRUMENTED_INSTRUCTION,

Python/opcode_targets.h

@@ -234,7 +234,6 @@ static void *opcode_targets[256] = {
     &&TARGET_UNPACK_SEQUENCE_TWO_TUPLE,
     &&_unknown_opcode,
     &&_unknown_opcode,
-    &&_unknown_opcode,
     &&TARGET_INSTRUMENTED_END_FOR,
     &&TARGET_INSTRUMENTED_POP_ITER,
     &&TARGET_INSTRUMENTED_END_SEND,
@@ -249,6 +248,7 @@ static void *opcode_targets[256] = {
     &&TARGET_INSTRUMENTED_RESUME,
     &&TARGET_INSTRUMENTED_RETURN_VALUE,
     &&TARGET_INSTRUMENTED_YIELD_VALUE,
+    &&TARGET_INSTRUMENTED_END_ASYNC_FOR,
     &&TARGET_INSTRUMENTED_LOAD_SUPER_ATTR,
     &&TARGET_INSTRUMENTED_CALL,
     &&TARGET_INSTRUMENTED_CALL_KW,
@@ -365,6 +365,7 @@ Py_PRESERVE_NONE_CC static PyObject *_TAIL_CALL_IMPORT_NAME(TAIL_CALL_PARAMS);
 Py_PRESERVE_NONE_CC static PyObject *_TAIL_CALL_INSTRUMENTED_CALL(TAIL_CALL_PARAMS);
 Py_PRESERVE_NONE_CC static PyObject *_TAIL_CALL_INSTRUMENTED_CALL_FUNCTION_EX(TAIL_CALL_PARAMS);
 Py_PRESERVE_NONE_CC static PyObject *_TAIL_CALL_INSTRUMENTED_CALL_KW(TAIL_CALL_PARAMS);
+Py_PRESERVE_NONE_CC static PyObject *_TAIL_CALL_INSTRUMENTED_END_ASYNC_FOR(TAIL_CALL_PARAMS);
 Py_PRESERVE_NONE_CC static PyObject *_TAIL_CALL_INSTRUMENTED_END_FOR(TAIL_CALL_PARAMS);
 Py_PRESERVE_NONE_CC static PyObject *_TAIL_CALL_INSTRUMENTED_END_SEND(TAIL_CALL_PARAMS);
 Py_PRESERVE_NONE_CC static PyObject *_TAIL_CALL_INSTRUMENTED_FOR_ITER(TAIL_CALL_PARAMS);
@@ -598,6 +599,7 @@ static py_tail_call_funcptr INSTRUCTION_TABLE[256] = {
     [INSTRUMENTED_CALL] = _TAIL_CALL_INSTRUMENTED_CALL,
     [INSTRUMENTED_CALL_FUNCTION_EX] = _TAIL_CALL_INSTRUMENTED_CALL_FUNCTION_EX,
     [INSTRUMENTED_CALL_KW] = _TAIL_CALL_INSTRUMENTED_CALL_KW,
+    [INSTRUMENTED_END_ASYNC_FOR] = _TAIL_CALL_INSTRUMENTED_END_ASYNC_FOR,
     [INSTRUMENTED_END_FOR] = _TAIL_CALL_INSTRUMENTED_END_FOR,
     [INSTRUMENTED_END_SEND] = _TAIL_CALL_INSTRUMENTED_END_SEND,
     [INSTRUMENTED_FOR_ITER] = _TAIL_CALL_INSTRUMENTED_FOR_ITER,
@@ -757,6 +759,5 @@ static py_tail_call_funcptr INSTRUCTION_TABLE[256] = {
     [148] = _TAIL_CALL_UNKNOWN_OPCODE,
     [232] = _TAIL_CALL_UNKNOWN_OPCODE,
     [233] = _TAIL_CALL_UNKNOWN_OPCODE,
-    [234] = _TAIL_CALL_UNKNOWN_OPCODE,
 };
 #endif /* Py_TAIL_CALL_INTERP */