gh-108654: restore comprehension locals before handling exception (#108659)

Co-authored-by: Dong-hee Na <donghee.na92@gmail.com>
This commit is contained in:
Carl Meyer 2023-08-30 17:50:50 -06:00 committed by GitHub
parent f59c66e8c8
commit d52c4482a8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 90 additions and 14 deletions

View file

@ -5529,6 +5529,8 @@ typedef struct {
PyObject *pushed_locals;
PyObject *temp_symbols;
PyObject *fast_hidden;
jump_target_label cleanup;
jump_target_label end;
} inlined_comprehension_state;
static int
@ -5639,11 +5641,45 @@ push_inlined_comprehension_state(struct compiler *c, location loc,
// `pushed_locals` on the stack, but this will be reversed when we swap
// out the comprehension result in pop_inlined_comprehension_state
ADDOP_I(c, loc, SWAP, PyList_GET_SIZE(state->pushed_locals) + 1);
// Add our own cleanup handler to restore comprehension locals in case
// of exception, so they have the correct values inside an exception
// handler or finally block.
NEW_JUMP_TARGET_LABEL(c, cleanup);
state->cleanup = cleanup;
NEW_JUMP_TARGET_LABEL(c, end);
state->end = end;
// no need to push an fblock for this "virtual" try/finally; there can't
// be return/continue/break inside a comprehension
ADDOP_JUMP(c, loc, SETUP_FINALLY, cleanup);
}
return SUCCESS;
}
static int
restore_inlined_comprehension_locals(struct compiler *c, location loc,
inlined_comprehension_state state)
{
PyObject *k;
// pop names we pushed to stack earlier
Py_ssize_t npops = PyList_GET_SIZE(state.pushed_locals);
// Preserve the comprehension result (or exception) as TOS. This
// reverses the SWAP we did in push_inlined_comprehension_state to get
// the outermost iterable to TOS, so we can still just iterate
// pushed_locals in simple reverse order
ADDOP_I(c, loc, SWAP, npops + 1);
for (Py_ssize_t i = npops - 1; i >= 0; --i) {
k = PyList_GetItem(state.pushed_locals, i);
if (k == NULL) {
return ERROR;
}
ADDOP_NAME(c, loc, STORE_FAST_MAYBE_NULL, k, varnames);
}
return SUCCESS;
}
static int
pop_inlined_comprehension_state(struct compiler *c, location loc,
inlined_comprehension_state state)
@ -5660,19 +5696,22 @@ pop_inlined_comprehension_state(struct compiler *c, location loc,
Py_CLEAR(state.temp_symbols);
}
if (state.pushed_locals) {
// pop names we pushed to stack earlier
Py_ssize_t npops = PyList_GET_SIZE(state.pushed_locals);
// Preserve the list/dict/set result of the comprehension as TOS. This
// reverses the SWAP we did in push_inlined_comprehension_state to get
// the outermost iterable to TOS, so we can still just iterate
// pushed_locals in simple reverse order
ADDOP_I(c, loc, SWAP, npops + 1);
for (Py_ssize_t i = npops - 1; i >= 0; --i) {
k = PyList_GetItem(state.pushed_locals, i);
if (k == NULL) {
return ERROR;
}
ADDOP_NAME(c, loc, STORE_FAST_MAYBE_NULL, k, varnames);
ADDOP(c, NO_LOCATION, POP_BLOCK);
ADDOP_JUMP(c, NO_LOCATION, JUMP, state.end);
// cleanup from an exception inside the comprehension
USE_LABEL(c, state.cleanup);
// discard incomplete comprehension result (beneath exc on stack)
ADDOP_I(c, NO_LOCATION, SWAP, 2);
ADDOP(c, NO_LOCATION, POP_TOP);
if (restore_inlined_comprehension_locals(c, loc, state) < 0) {
return ERROR;
}
ADDOP_I(c, NO_LOCATION, RERAISE, 0);
USE_LABEL(c, state.end);
if (restore_inlined_comprehension_locals(c, loc, state) < 0) {
return ERROR;
}
Py_CLEAR(state.pushed_locals);
}
@ -5715,7 +5754,7 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,
expr_ty val)
{
PyCodeObject *co = NULL;
inlined_comprehension_state inline_state = {NULL, NULL};
inlined_comprehension_state inline_state = {NULL, NULL, NULL, NO_LABEL, NO_LABEL};
comprehension_ty outermost;
int scope_type = c->u->u_scope_type;
int is_top_level_await = IS_TOP_LEVEL_AWAIT(c);