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
This commit is contained in:
Mark Shannon 2025-03-07 14:30:31 +00:00 committed by GitHub
parent e5527f2cdd
commit 89df62c120
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 235 additions and 154 deletions

View file

@ -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,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;
}

View file

@ -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;

View file

@ -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;
}

View file

@ -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",

View file

@ -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);

View file

@ -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;
}

View file

@ -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,