Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Aug 15, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
165 changes: 102 additions & 63 deletions Python/compile.c
Original file line number Diff line number Diff line change
Expand Up @@ -5371,38 +5371,28 @@ typedef struct {
PyObject *temp_symbols;
PyObject *fast_hidden;
jump_target_label cleanup;
jump_target_label end;
} inlined_comprehension_state;

static int
push_inlined_comprehension_state(struct compiler *c, location loc,
PySTEntryObject *entry,
inlined_comprehension_state *state)
compiler_tweak_inlined_comprehension_scopes(struct compiler *c, location loc,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure what this does, maybe change "tweak" to something more informative?

Can we also add a comment as we why we need to "tweak" the "state", rather than creating a new symbol table entry for the comprehension?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have this issue to look into this area - let's discuss there. #122985

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(Sorry, I already clicked auto-merge before your comment came through).

PySTEntryObject *entry,
inlined_comprehension_state *state)
{
int in_class_block = (SYMTABLE_ENTRY(c)->ste_type == ClassBlock) && !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;
Py_ssize_t pos = 0;
while (PyDict_Next(entry->ste_symbols, &pos, &k, &v)) {
long symbol = PyLong_AsLong(v);
if (symbol == -1 && PyErr_Occurred()) {
return ERROR;
}
assert(symbol >= 0 || PyErr_Occurred());
RETURN_IF_ERROR(symbol);
long scope = SYMBOL_TO_SCOPE(symbol);
PyObject *outv = PyDict_GetItemWithError(SYMTABLE_ENTRY(c)->ste_symbols, k);
if (outv == NULL) {
if (PyErr_Occurred()) {
return ERROR;
}
outv = _PyLong_GetZero();
}
long outsymbol = PyLong_AsLong(outv);
if (outsymbol == -1 && PyErr_Occurred()) {
return ERROR;
}

long outsymbol = _PyST_GetSymbol(SYMTABLE_ENTRY(c), k);
RETURN_IF_ERROR(outsymbol);
long outsc = SYMBOL_TO_SCOPE(outsymbol);

// If a name has different scope inside than outside the comprehension,
// we need to temporarily handle it with the right scope while
// compiling the comprehension. If it's free in the comprehension
Expand All @@ -5422,16 +5412,16 @@ push_inlined_comprehension_state(struct compiler *c, location loc,
// update the symbol to the in-comprehension version and save
// the outer version; we'll restore it after running the
// comprehension
Py_INCREF(outv);
if (PyDict_SetItem(SYMTABLE_ENTRY(c)->ste_symbols, k, v) < 0) {
Py_DECREF(outv);
return ERROR;
}
if (PyDict_SetItem(state->temp_symbols, k, outv) < 0) {
Py_DECREF(outv);
PyObject *outv = PyLong_FromLong(outsymbol);
if (outv == NULL) {
return ERROR;
}
int res = PyDict_SetItem(state->temp_symbols, k, outv);
Py_DECREF(outv);
RETURN_IF_ERROR(res);
}
// locals handling for names bound in comprehension (DEF_LOCAL |
// DEF_NONLOCAL occurs in assignment expression to nonlocal)
Expand All @@ -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) {
return ERROR;
}
int orig_is_true = (orig == Py_True);
Py_XDECREF(orig);
if (!orig_is_true) {
assert(orig == NULL || orig == Py_True || orig == Py_False);
if (orig != Py_True) {
if (PyDict_SetItem(c->u->u_metadata.u_fasthidden, k, Py_True) < 0) {
return ERROR;
}
Expand All @@ -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
// outer scope; push existing value (which may be NULL if
// not defined) on stack
Expand Down Expand Up @@ -5496,31 +5512,40 @@ push_inlined_comprehension_state(struct compiler *c, location loc,
// 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
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;
}

static int
restore_inlined_comprehension_locals(struct compiler *c, location loc,
inlined_comprehension_state state)
inlined_comprehension_state *state)
{
PyObject *k;
// 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
// reverses the SWAP we did in push_inlined_comprehension_state to get
// the outermost iterable to TOS, so we can still just iterate
// 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);
k = PyList_GetItem(state->pushed_locals, i);
if (k == NULL) {
return ERROR;
}
Expand All @@ -5530,43 +5555,47 @@ restore_inlined_comprehension_locals(struct compiler *c, location loc,
}

static int
pop_inlined_comprehension_state(struct compiler *c, location loc,
inlined_comprehension_state state)
codegen_pop_inlined_comprehension_locals(struct compiler *c, location loc,
inlined_comprehension_state *state)
{
c->u->u_in_inlined_comp--;
PyObject *k, *v;
Py_ssize_t pos = 0;
if (state.temp_symbols) {
while (PyDict_Next(state.temp_symbols, &pos, &k, &v)) {
if (PyDict_SetItem(SYMTABLE_ENTRY(c)->ste_symbols, k, v)) {
return ERROR;
}
}
Py_CLEAR(state.temp_symbols);
}
if (state.pushed_locals) {
if (state->pushed_locals) {
ADDOP(c, NO_LOCATION, POP_BLOCK);
ADDOP_JUMP(c, NO_LOCATION, JUMP_NO_INTERRUPT, state.end);

NEW_JUMP_TARGET_LABEL(c, end);
ADDOP_JUMP(c, NO_LOCATION, JUMP_NO_INTERRUPT, end);

// cleanup from an exception inside the comprehension
USE_LABEL(c, state.cleanup);
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;
}
RETURN_IF_ERROR(restore_inlined_comprehension_locals(c, loc, state));
ADDOP_I(c, NO_LOCATION, RERAISE, 0);

USE_LABEL(c, state.end);
if (restore_inlined_comprehension_locals(c, loc, state) < 0) {
return ERROR;
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)) {
return ERROR;
}
}
Py_CLEAR(state.pushed_locals);
Py_CLEAR(state->temp_symbols);
}
if (state.fast_hidden) {
while (PySet_Size(state.fast_hidden) > 0) {
PyObject *k = PySet_Pop(state.fast_hidden);
if (state->fast_hidden) {
while (PySet_Size(state->fast_hidden) > 0) {
PyObject *k = PySet_Pop(state->fast_hidden);
if (k == NULL) {
return ERROR;
}
Expand All @@ -5578,11 +5607,21 @@ pop_inlined_comprehension_state(struct compiler *c, location loc,
}
Py_DECREF(k);
}
Py_CLEAR(state.fast_hidden);
Py_CLEAR(state->fast_hidden);
}
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
compiler_comprehension_iter(struct compiler *c, location loc,
comprehension_ty comp)
Expand All @@ -5603,7 +5642,7 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,
expr_ty val)
{
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;
#ifndef NDEBUG
int scope_type = c->u->u_scope_type;
Expand Down Expand Up @@ -5671,7 +5710,7 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,
}

if (is_inlined) {
if (pop_inlined_comprehension_state(c, loc, inline_state)) {
if (pop_inlined_comprehension_state(c, loc, &inline_state)) {
goto error;
}
return SUCCESS;
Expand Down
Loading