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

5 files changed, 70 insertions, 3 deletions

diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index 9f1309580a6..dc709a7c053 100644
--- a/Python/bytecodes.c
+++ b/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
diff --git a/Python/codegen.c b/Python/codegen.c
index 7669cc5a65b..ecad8c22bdf 100644
--- a/Python/codegen.c
+++ b/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);
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
index 311faf340fc..5f1198699c4 100644
--- a/Python/generated_cases.c.h
+++ b/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;
diff --git a/Python/instrumentation.c b/Python/instrumentation.c
index 12fd35f4356..4e7ca808b3c 100644
--- a/Python/instrumentation.c
+++ b/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,
diff --git a/Python/opcode_targets.h b/Python/opcode_targets.h
index 142fbc33641..0435d0841db 100644
--- a/Python/opcode_targets.h
+++ b/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 */