[3.12] gh-108654: restore comprehension locals before handling exception (GH-108659) (#108700)

gh-108654: restore comprehension locals before handling exception (GH-108659)
(cherry picked from commit d52c4482a8)

Co-authored-by: Carl Meyer <carl@oddbird.net>
Co-authored-by: Dong-hee Na <donghee.na92@gmail.com>
This commit is contained in:
Miss Islington (bot) 2023-08-31 06:32:17 -07:00 committed by GitHub
parent 320d398262
commit 1a15d20b75
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 90 additions and 14 deletions

View file

@ -5438,6 +5438,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
@ -5543,11 +5545,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)
@ -5564,19 +5600,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);
}
@ -5619,7 +5658,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);