bpo-43693: Eliminate unused "fast locals". (gh-26587)

Currently, if an arg value escapes (into the closure for an inner function) we end up allocating two indices in the fast locals even though only one gets used.  Additionally, using the lower index would be better in some cases, such as with no-arg `super()`.  To address this, we update the compiler to fix the offsets so each variable only gets one "fast local".  As a consequence, now some cell offsets are interspersed with the locals (only when an arg escapes to an inner function).

https://bugs.python.org/issue43693

1 files changed, 36 insertions, 74 deletions

diff --git a/Objects/codeobject.c b/Objects/codeobject.c
index e054c43..11a83d4 100644
--- a/Objects/codeobject.c
+++ b/Objects/codeobject.c
@@ -162,43 +162,28 @@ _Py_set_localsplus_info(int offset, PyObject *name, _PyLocalsPlusKind kind,
     Py_INCREF(name);
     PyTuple_SET_ITEM(names, offset, name);
     kinds[offset] = kind;
-
-    if (kind == CO_FAST_CELL) {
-        // Cells can overlap with args, so mark those cases.
-        int nlocalsplus = (int)PyTuple_GET_SIZE(names);
-        for (int i = 0; i < nlocalsplus; i++) {
-            _PyLocalsPlusKind kind = kinds[i];
-            if (kind && !(kind & CO_FAST_LOCAL)) {
-                // We've moved past the locals.
-                break;
-            }
-            PyObject *varname = PyTuple_GET_ITEM(names, i);
-            int cmp = PyUnicode_Compare(name, varname);
-            if (cmp == 0) {
-                kinds[i] |= CO_FAST_CELL;
-                break;
-            }
-            assert(cmp > 0 || !PyErr_Occurred());
-        }
-    }
 }
 
 static void
 get_localsplus_counts(PyObject *names, _PyLocalsPlusKinds kinds,
-                      int *pnlocals, int *pncellvars,
+                      int *pnlocals, int *pnplaincellvars, int *pncellvars,
                       int *pnfreevars)
 {
     int nlocals = 0;
+    int nplaincellvars = 0;
     int ncellvars = 0;
     int nfreevars = 0;
-    int nlocalsplus = Py_SAFE_DOWNCAST(PyTuple_GET_SIZE(names),
-                                       Py_ssize_t, int);
+    Py_ssize_t nlocalsplus = PyTuple_GET_SIZE(names);
     for (int i = 0; i < nlocalsplus; i++) {
         if (kinds[i] & CO_FAST_LOCAL) {
             nlocals += 1;
+            if (kinds[i] & CO_FAST_CELL) {
+                ncellvars += 1;
+            }
         }
         else if (kinds[i] & CO_FAST_CELL) {
             ncellvars += 1;
+            nplaincellvars += 1;
         }
         else if (kinds[i] & CO_FAST_FREE) {
             nfreevars += 1;
@@ -207,6 +192,9 @@ get_localsplus_counts(PyObject *names, _PyLocalsPlusKinds kinds,
     if (pnlocals != NULL) {
         *pnlocals = nlocals;
     }
+    if (pnplaincellvars != NULL) {
+        *pnplaincellvars = nplaincellvars;
+    }
     if (pncellvars != NULL) {
         *pncellvars = ncellvars;
     }
@@ -227,10 +215,6 @@ get_localsplus_names(PyCodeObject *co, _PyLocalsPlusKind kind, int num)
         if ((co->co_localspluskinds[offset] & kind) == 0) {
             continue;
         }
-        // For now there may be duplicates, which we ignore.
-        if (kind == CO_FAST_CELL && co->co_localspluskinds[offset] != kind) {
-            continue;
-        }
         assert(index < num);
         PyObject *name = PyTuple_GET_ITEM(co->co_localsplusnames, offset);
         Py_INCREF(name);
@@ -283,7 +267,7 @@ _PyCode_Validate(struct _PyCodeConstructor *con)
      * here to avoid the possibility of overflow (however remote). */
     int nlocals;
     get_localsplus_counts(con->localsplusnames, con->localspluskinds,
-                          &nlocals, NULL, NULL);
+                          &nlocals, NULL, NULL, NULL);
     int nplainlocals = nlocals -
                        con->argcount -
                        con->kwonlyargcount -
@@ -301,9 +285,9 @@ static void
 init_code(PyCodeObject *co, struct _PyCodeConstructor *con)
 {
     int nlocalsplus = (int)PyTuple_GET_SIZE(con->localsplusnames);
-    int nlocals, ncellvars, nfreevars;
+    int nlocals, nplaincellvars, ncellvars, nfreevars;
     get_localsplus_counts(con->localsplusnames, con->localspluskinds,
-                          &nlocals, &ncellvars, &nfreevars);
+                          &nlocals, &nplaincellvars, &ncellvars, &nfreevars);
 
     Py_INCREF(con->filename);
     co->co_filename = con->filename;
@@ -338,9 +322,9 @@ init_code(PyCodeObject *co, struct _PyCodeConstructor *con)
     co->co_exceptiontable = con->exceptiontable;
 
     /* derived values */
-    co->co_cell2arg = NULL;  // This will be set soon.
     co->co_nlocalsplus = nlocalsplus;
     co->co_nlocals = nlocals;
+    co->co_nplaincellvars = nplaincellvars;
     co->co_ncellvars = ncellvars;
     co->co_nfreevars = nfreevars;
     co->co_varnames = NULL;
@@ -392,44 +376,6 @@ _PyCode_New(struct _PyCodeConstructor *con)
         co->co_flags &= ~CO_NOFREE;
     }
 
-    /* Create mapping between cells and arguments if needed. */
-    if (co->co_ncellvars) {
-        int totalargs = co->co_argcount +
-                        co->co_kwonlyargcount +
-                        ((co->co_flags & CO_VARARGS) != 0) +
-                        ((co->co_flags & CO_VARKEYWORDS) != 0);
-        assert(totalargs <= co->co_nlocals);
-        /* Find cells which are also arguments. */
-        for (int i = 0; i < co->co_ncellvars; i++) {
-            PyObject *cellname = PyTuple_GET_ITEM(co->co_localsplusnames,
-                                                  i + co->co_nlocals);
-            for (int j = 0; j < totalargs; j++) {
-                PyObject *argname = PyTuple_GET_ITEM(co->co_localsplusnames, j);
-                int cmp = PyUnicode_Compare(cellname, argname);
-                if (cmp == -1 && PyErr_Occurred()) {
-                    Py_DECREF(co);
-                    return NULL;
-                }
-                if (cmp == 0) {
-                    if (co->co_cell2arg == NULL) {
-                        co->co_cell2arg = PyMem_NEW(int, co->co_ncellvars);
-                        if (co->co_cell2arg == NULL) {
-                            Py_DECREF(co);
-                            PyErr_NoMemory();
-                            return NULL;
-                        }
-                        for (int k = 0; k < co->co_ncellvars; k++) {
-                            co->co_cell2arg[k] = CO_CELL_NOT_AN_ARG;
-                        }
-                    }
-                    co->co_cell2arg[i] = j;
-                    // Go to the next cell name.
-                    break;
-                }
-            }
-        }
-    }
-
     return co;
 }
 
@@ -478,6 +424,23 @@ PyCode_NewWithPosOnlyArgs(int argcount, int posonlyargcount, int kwonlyargcount,
     }
     for (int i = 0; i < ncellvars; i++, offset++) {
         PyObject *name = PyTuple_GET_ITEM(cellvars, i);
+        int argoffset = -1;
+        for (int j = 0; j < nvarnames; j++) {
+            int cmp = PyUnicode_Compare(PyTuple_GET_ITEM(varnames, j),
+                                        name);
+            assert(!PyErr_Occurred());
+            if (cmp == 0) {
+                argoffset = j;
+                break;
+            }
+        }
+        if (argoffset >= 0) {
+            // Merge the localsplus indices.
+            nlocalsplus -= 1;
+            offset -= 1;
+            localspluskinds[argoffset] |= CO_FAST_CELL;
+            continue;
+        }
         _Py_set_localsplus_info(offset, name, CO_FAST_CELL,
                                localsplusnames, localspluskinds);
     }
@@ -486,6 +449,11 @@ PyCode_NewWithPosOnlyArgs(int argcount, int posonlyargcount, int kwonlyargcount,
         _Py_set_localsplus_info(offset, name, CO_FAST_FREE,
                                localsplusnames, localspluskinds);
     }
+    // If any cells were args then nlocalsplus will have shrunk.
+    // We don't bother resizing localspluskinds.
+    if (_PyTuple_Resize(&localsplusnames, nlocalsplus) < 0) {
+        goto error;
+    }
 
     struct _PyCodeConstructor con = {
         .filename = filename,
@@ -1182,8 +1150,6 @@ code_dealloc(PyCodeObject *co)
     Py_XDECREF(co->co_name);
     Py_XDECREF(co->co_linetable);
     Py_XDECREF(co->co_exceptiontable);
-    if (co->co_cell2arg != NULL)
-        PyMem_Free(co->co_cell2arg);
     if (co->co_weakreflist != NULL)
         PyObject_ClearWeakRefs((PyObject*)co);
     if (co->co_quickened) {
@@ -1377,10 +1343,6 @@ code_sizeof(PyCodeObject *co, PyObject *Py_UNUSED(args))
                (co_extra->ce_size-1) * sizeof(co_extra->ce_extras[0]);
     }
 
-    if (co->co_cell2arg != NULL && co->co_cellvars != NULL) {
-        res += co->co_ncellvars * sizeof(Py_ssize_t);
-    }
-
     if (co->co_quickened != NULL) {
         Py_ssize_t count = co->co_quickened[0].entry.zero.cache_count;
         count += (PyBytes_GET_SIZE(co->co_code)+sizeof(SpecializedCacheEntry)-1)/