gh-121404: split compiler's push/pop_inlined_comprehension_state into codegen and compiler parts (#123021)

This commit is contained in:
Irit Katriel 2024-08-15 12:08:54 +01:00 committed by GitHub
parent 3950392f48
commit b6cb435ac0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 107 additions and 68 deletions

View File

@ -5371,38 +5371,28 @@ typedef struct {
PyObject *temp_symbols; PyObject *temp_symbols;
PyObject *fast_hidden; PyObject *fast_hidden;
jump_target_label cleanup; jump_target_label cleanup;
jump_target_label end;
} inlined_comprehension_state; } inlined_comprehension_state;
static int static int
push_inlined_comprehension_state(struct compiler *c, location loc, compiler_tweak_inlined_comprehension_scopes(struct compiler *c, location loc,
PySTEntryObject *entry, PySTEntryObject *entry,
inlined_comprehension_state *state) inlined_comprehension_state *state)
{ {
int in_class_block = (SYMTABLE_ENTRY(c)->ste_type == ClassBlock) && !c->u->u_in_inlined_comp; int in_class_block = (SYMTABLE_ENTRY(c)->ste_type == ClassBlock) && !c->u->u_in_inlined_comp;
c->u->u_in_inlined_comp++; c->u->u_in_inlined_comp++;
// iterate over names bound in the comprehension and ensure we isolate
// them from the outer scope as needed
PyObject *k, *v; PyObject *k, *v;
Py_ssize_t pos = 0; Py_ssize_t pos = 0;
while (PyDict_Next(entry->ste_symbols, &pos, &k, &v)) { while (PyDict_Next(entry->ste_symbols, &pos, &k, &v)) {
long symbol = PyLong_AsLong(v); long symbol = PyLong_AsLong(v);
if (symbol == -1 && PyErr_Occurred()) { assert(symbol >= 0 || PyErr_Occurred());
return ERROR; RETURN_IF_ERROR(symbol);
}
long scope = SYMBOL_TO_SCOPE(symbol); long scope = SYMBOL_TO_SCOPE(symbol);
PyObject *outv = PyDict_GetItemWithError(SYMTABLE_ENTRY(c)->ste_symbols, k);
if (outv == NULL) { long outsymbol = _PyST_GetSymbol(SYMTABLE_ENTRY(c), k);
if (PyErr_Occurred()) { RETURN_IF_ERROR(outsymbol);
return ERROR;
}
outv = _PyLong_GetZero();
}
long outsymbol = PyLong_AsLong(outv);
if (outsymbol == -1 && PyErr_Occurred()) {
return ERROR;
}
long outsc = SYMBOL_TO_SCOPE(outsymbol); long outsc = SYMBOL_TO_SCOPE(outsymbol);
// If a name has different scope inside than outside the comprehension, // If a name has different scope inside than outside the comprehension,
// we need to temporarily handle it with the right scope while // we need to temporarily handle it with the right scope while
// compiling the comprehension. If it's free in the comprehension // compiling the comprehension. If it's free in the comprehension
@ -5422,16 +5412,16 @@ push_inlined_comprehension_state(struct compiler *c, location loc,
// update the symbol to the in-comprehension version and save // update the symbol to the in-comprehension version and save
// the outer version; we'll restore it after running the // the outer version; we'll restore it after running the
// comprehension // comprehension
Py_INCREF(outv);
if (PyDict_SetItem(SYMTABLE_ENTRY(c)->ste_symbols, k, v) < 0) { if (PyDict_SetItem(SYMTABLE_ENTRY(c)->ste_symbols, k, v) < 0) {
Py_DECREF(outv);
return ERROR; return ERROR;
} }
if (PyDict_SetItem(state->temp_symbols, k, outv) < 0) { PyObject *outv = PyLong_FromLong(outsymbol);
Py_DECREF(outv); if (outv == NULL) {
return ERROR; return ERROR;
} }
int res = PyDict_SetItem(state->temp_symbols, k, outv);
Py_DECREF(outv); Py_DECREF(outv);
RETURN_IF_ERROR(res);
} }
// locals handling for names bound in comprehension (DEF_LOCAL | // locals handling for names bound in comprehension (DEF_LOCAL |
// DEF_NONLOCAL occurs in assignment expression to nonlocal) // DEF_NONLOCAL occurs in assignment expression to nonlocal)
@ -5442,9 +5432,8 @@ push_inlined_comprehension_state(struct compiler *c, location loc,
if (PyDict_GetItemRef(c->u->u_metadata.u_fasthidden, k, &orig) < 0) { if (PyDict_GetItemRef(c->u->u_metadata.u_fasthidden, k, &orig) < 0) {
return ERROR; return ERROR;
} }
int orig_is_true = (orig == Py_True); assert(orig == NULL || orig == Py_True || orig == Py_False);
Py_XDECREF(orig); if (orig != Py_True) {
if (!orig_is_true) {
if (PyDict_SetItem(c->u->u_metadata.u_fasthidden, k, Py_True) < 0) { if (PyDict_SetItem(c->u->u_metadata.u_fasthidden, k, Py_True) < 0) {
return ERROR; return ERROR;
} }
@ -5459,6 +5448,33 @@ push_inlined_comprehension_state(struct compiler *c, location loc,
} }
} }
} }
}
}
return SUCCESS;
}
static int
codegen_push_inlined_comprehension_locals(struct compiler *c, location loc,
PySTEntryObject *comp,
inlined_comprehension_state *state)
{
int in_class_block = (SYMTABLE_ENTRY(c)->ste_type == ClassBlock) && !c->u->u_in_inlined_comp;
PySTEntryObject *outer = SYMTABLE_ENTRY(c);
// iterate over names bound in the comprehension and ensure we isolate
// them from the outer scope as needed
PyObject *k, *v;
Py_ssize_t pos = 0;
while (PyDict_Next(comp->ste_symbols, &pos, &k, &v)) {
long symbol = PyLong_AsLong(v);
assert(symbol >= 0 || PyErr_Occurred());
RETURN_IF_ERROR(symbol);
long scope = SYMBOL_TO_SCOPE(symbol);
long outsymbol = _PyST_GetSymbol(outer, k);
RETURN_IF_ERROR(outsymbol);
long outsc = SYMBOL_TO_SCOPE(outsymbol);
if ((symbol & DEF_LOCAL && !(symbol & DEF_NONLOCAL)) || in_class_block) {
// local names bound in comprehension must be isolated from // local names bound in comprehension must be isolated from
// outer scope; push existing value (which may be NULL if // outer scope; push existing value (which may be NULL if
// not defined) on stack // not defined) on stack
@ -5496,31 +5512,40 @@ push_inlined_comprehension_state(struct compiler *c, location loc,
// handler or finally block. // handler or finally block.
NEW_JUMP_TARGET_LABEL(c, cleanup); NEW_JUMP_TARGET_LABEL(c, cleanup);
state->cleanup = 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 // no need to push an fblock for this "virtual" try/finally; there can't
// be return/continue/break inside a comprehension // be return/continue/break inside a comprehension
ADDOP_JUMP(c, loc, SETUP_FINALLY, cleanup); ADDOP_JUMP(c, loc, SETUP_FINALLY, cleanup);
} }
return SUCCESS;
}
static int
push_inlined_comprehension_state(struct compiler *c, location loc,
PySTEntryObject *comp,
inlined_comprehension_state *state)
{
RETURN_IF_ERROR(
compiler_tweak_inlined_comprehension_scopes(c, loc, comp, state));
RETURN_IF_ERROR(
codegen_push_inlined_comprehension_locals(c, loc, comp, state));
return SUCCESS; return SUCCESS;
} }
static int static int
restore_inlined_comprehension_locals(struct compiler *c, location loc, restore_inlined_comprehension_locals(struct compiler *c, location loc,
inlined_comprehension_state state) inlined_comprehension_state *state)
{ {
PyObject *k; PyObject *k;
// pop names we pushed to stack earlier // pop names we pushed to stack earlier
Py_ssize_t npops = PyList_GET_SIZE(state.pushed_locals); Py_ssize_t npops = PyList_GET_SIZE(state->pushed_locals);
// Preserve the comprehension result (or exception) as TOS. This // Preserve the comprehension result (or exception) as TOS. This
// reverses the SWAP we did in push_inlined_comprehension_state to get // reverses the SWAP we did in push_inlined_comprehension_state
// the outermost iterable to TOS, so we can still just iterate // to get the outermost iterable to TOS, so we can still just iterate
// pushed_locals in simple reverse order // pushed_locals in simple reverse order
ADDOP_I(c, loc, SWAP, npops + 1); ADDOP_I(c, loc, SWAP, npops + 1);
for (Py_ssize_t i = npops - 1; i >= 0; --i) { for (Py_ssize_t i = npops - 1; i >= 0; --i) {
k = PyList_GetItem(state.pushed_locals, i); k = PyList_GetItem(state->pushed_locals, i);
if (k == NULL) { if (k == NULL) {
return ERROR; return ERROR;
} }
@ -5530,43 +5555,47 @@ restore_inlined_comprehension_locals(struct compiler *c, location loc,
} }
static int static int
pop_inlined_comprehension_state(struct compiler *c, location loc, codegen_pop_inlined_comprehension_locals(struct compiler *c, location loc,
inlined_comprehension_state state) inlined_comprehension_state *state)
{ {
c->u->u_in_inlined_comp--; if (state->pushed_locals) {
PyObject *k, *v; ADDOP(c, NO_LOCATION, POP_BLOCK);
Py_ssize_t pos = 0;
if (state.temp_symbols) { NEW_JUMP_TARGET_LABEL(c, end);
while (PyDict_Next(state.temp_symbols, &pos, &k, &v)) { ADDOP_JUMP(c, NO_LOCATION, JUMP_NO_INTERRUPT, 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);
RETURN_IF_ERROR(restore_inlined_comprehension_locals(c, loc, state));
ADDOP_I(c, NO_LOCATION, RERAISE, 0);
USE_LABEL(c, end);
RETURN_IF_ERROR(restore_inlined_comprehension_locals(c, loc, state));
Py_CLEAR(state->pushed_locals);
}
return SUCCESS;
}
static int
compiler_revert_inlined_comprehension_scopes(struct compiler *c, location loc,
inlined_comprehension_state *state)
{
if (state->temp_symbols) {
PyObject *k, *v;
Py_ssize_t pos = 0;
while (PyDict_Next(state->temp_symbols, &pos, &k, &v)) {
if (PyDict_SetItem(SYMTABLE_ENTRY(c)->ste_symbols, k, v)) { if (PyDict_SetItem(SYMTABLE_ENTRY(c)->ste_symbols, k, v)) {
return ERROR; return ERROR;
} }
} }
Py_CLEAR(state.temp_symbols); Py_CLEAR(state->temp_symbols);
} }
if (state.pushed_locals) { if (state->fast_hidden) {
ADDOP(c, NO_LOCATION, POP_BLOCK); while (PySet_Size(state->fast_hidden) > 0) {
ADDOP_JUMP(c, NO_LOCATION, JUMP_NO_INTERRUPT, state.end); PyObject *k = PySet_Pop(state->fast_hidden);
// 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);
}
if (state.fast_hidden) {
while (PySet_Size(state.fast_hidden) > 0) {
PyObject *k = PySet_Pop(state.fast_hidden);
if (k == NULL) { if (k == NULL) {
return ERROR; return ERROR;
} }
@ -5578,11 +5607,21 @@ pop_inlined_comprehension_state(struct compiler *c, location loc,
} }
Py_DECREF(k); Py_DECREF(k);
} }
Py_CLEAR(state.fast_hidden); Py_CLEAR(state->fast_hidden);
} }
return SUCCESS; return SUCCESS;
} }
static int
pop_inlined_comprehension_state(struct compiler *c, location loc,
inlined_comprehension_state *state)
{
c->u->u_in_inlined_comp--;
RETURN_IF_ERROR(codegen_pop_inlined_comprehension_locals(c, loc, state));
RETURN_IF_ERROR(compiler_revert_inlined_comprehension_scopes(c, loc, state));
return SUCCESS;
}
static inline int static inline int
compiler_comprehension_iter(struct compiler *c, location loc, compiler_comprehension_iter(struct compiler *c, location loc,
comprehension_ty comp) comprehension_ty comp)
@ -5603,7 +5642,7 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,
expr_ty val) expr_ty val)
{ {
PyCodeObject *co = NULL; PyCodeObject *co = NULL;
inlined_comprehension_state inline_state = {NULL, NULL, NULL, NO_LABEL, NO_LABEL}; inlined_comprehension_state inline_state = {NULL, NULL, NULL, NO_LABEL};
comprehension_ty outermost; comprehension_ty outermost;
#ifndef NDEBUG #ifndef NDEBUG
int scope_type = c->u->u_scope_type; int scope_type = c->u->u_scope_type;
@ -5671,7 +5710,7 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,
} }
if (is_inlined) { if (is_inlined) {
if (pop_inlined_comprehension_state(c, loc, inline_state)) { if (pop_inlined_comprehension_state(c, loc, &inline_state)) {
goto error; goto error;
} }
return SUCCESS; return SUCCESS;