gh-87859: Track Code Object Local Kinds For Arguments (gh-132980)

Doing this was always the intention. I was finally motivated to find the time to do it.

See #87859 (comment).
This commit is contained in:
Eric Snow 2025-04-28 20:21:47 -06:00 committed by GitHub
parent 96a7fb93a8
commit 219d8d24b5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 372 additions and 194 deletions

View file

@ -482,33 +482,52 @@ extern void _Py_set_localsplus_info(int, PyObject *, unsigned char,
static int
compute_localsplus_info(_PyCompile_CodeUnitMetadata *umd, int nlocalsplus,
PyObject *names, PyObject *kinds)
int flags, PyObject *names, PyObject *kinds)
{
PyObject *k, *v;
Py_ssize_t pos = 0;
while (PyDict_Next(umd->u_varnames, &pos, &k, &v)) {
int offset = PyLong_AsInt(v);
if (offset == -1 && PyErr_Occurred()) {
return ERROR;
}
assert(offset >= 0);
assert(offset < nlocalsplus);
// For now we do not distinguish arg kinds.
_PyLocals_Kind kind = CO_FAST_LOCAL;
int has_key = PyDict_Contains(umd->u_fasthidden, k);
RETURN_IF_ERROR(has_key);
if (has_key) {
kind |= CO_FAST_HIDDEN;
}
// Set the locals kinds. Arg vars fill the first portion of the list.
struct {
int count;
_PyLocals_Kind kind;
} argvarkinds[6] = {
{(int)umd->u_posonlyargcount, CO_FAST_ARG_POS},
{(int)umd->u_argcount, CO_FAST_ARG_POS | CO_FAST_ARG_KW},
{(int)umd->u_kwonlyargcount, CO_FAST_ARG_KW},
{!!(flags & CO_VARARGS), CO_FAST_ARG_VAR | CO_FAST_ARG_POS},
{!!(flags & CO_VARKEYWORDS), CO_FAST_ARG_VAR | CO_FAST_ARG_KW},
{-1, 0}, // the remaining local vars
};
int max = 0;
for (int i = 0; i < 6; i++) {
max = argvarkinds[i].count < 0
? INT_MAX
: max + argvarkinds[i].count;
while (pos < max && PyDict_Next(umd->u_varnames, &pos, &k, &v)) {
int offset = PyLong_AsInt(v);
if (offset == -1 && PyErr_Occurred()) {
return ERROR;
}
assert(offset >= 0);
assert(offset < nlocalsplus);
has_key = PyDict_Contains(umd->u_cellvars, k);
RETURN_IF_ERROR(has_key);
if (has_key) {
kind |= CO_FAST_CELL;
}
_PyLocals_Kind kind = CO_FAST_LOCAL | argvarkinds[i].kind;
_Py_set_localsplus_info(offset, k, kind, names, kinds);
int has_key = PyDict_Contains(umd->u_fasthidden, k);
RETURN_IF_ERROR(has_key);
if (has_key) {
kind |= CO_FAST_HIDDEN;
}
has_key = PyDict_Contains(umd->u_cellvars, k);
RETURN_IF_ERROR(has_key);
if (has_key) {
kind |= CO_FAST_CELL;
}
_Py_set_localsplus_info(offset, k, kind, names, kinds);
}
}
int nlocals = (int)PyDict_GET_SIZE(umd->u_varnames);
@ -594,8 +613,10 @@ makecode(_PyCompile_CodeUnitMetadata *umd, struct assembler *a, PyObject *const_
if (localspluskinds == NULL) {
goto error;
}
if (compute_localsplus_info(umd, nlocalsplus,
localsplusnames, localspluskinds) == ERROR) {
if (compute_localsplus_info(
umd, nlocalsplus, code_flags,
localsplusnames, localspluskinds) == ERROR)
{
goto error;
}