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
fc3ca199dcae375c63fadb2b1023173dcb16ea60
commit fc3ca199dcae375c63fadb2b1023173dcb16ea60
diff --git a/Include/cpython/code.h b/Include/cpython/code.h
@@ -83,11 +83,11 @@ struct PyCodeObject {
 
     /* These fields are set with computed values on new code objects. */
 
-    int *co_cell2arg;           /* Maps cell vars which are arguments. */
     // redundant values (derived from co_localsplusnames and co_localspluskinds)
     int co_nlocalsplus;         /* number of local + cell + free variables */
     int co_nlocals;             /* number of local variables */
-    int co_ncellvars;           /* number of cell variables */
+    int co_nplaincellvars;      /* number of non-arg cell variables */
+    int co_ncellvars;           /* total number of cell variables */
     int co_nfreevars;           /* number of free variables */
     // lazily-computed values
     PyObject *co_varnames;      /* tuple of strings (local variable names) */
@@ -142,10 +142,6 @@ struct PyCodeObject {
 #define CO_FUTURE_GENERATOR_STOP  0x800000
 #define CO_FUTURE_ANNOTATIONS    0x1000000
 
-/* This value is found in the co_cell2arg array when the associated cell
-   variable does not correspond to an argument. */
-#define CO_CELL_NOT_AN_ARG (-1)
-
 /* This should be defined if a future statement modifies the syntax.
    For example, when a keyword is added.
 */

diff --git a/Include/internal/pycore_frame.h b/Include/internal/pycore_frame.h
@@ -32,8 +32,6 @@ _PyFrame_GetBuiltins(PyFrameObject *f)
 
 int _PyFrame_TakeLocals(PyFrameObject *f);
 
-PyAPI_FUNC(int) _PyFrame_OpAlreadyRan(PyFrameObject *f, int opcode, int oparg);
-
 #ifdef __cplusplus
 }
 #endif

diff --git a/Lib/importlib/_bootstrap_external.py b/Lib/importlib/_bootstrap_external.py
@@ -358,6 +358,7 @@ def _write_atomic(path, data, mode=0o666):
 #     Python 3.11a1 3453 (add co_fastlocalnames and co_fastlocalkinds)
 #     Python 3.11a1 3454 (compute cell offsets relative to locals bpo-43693)
 #     Python 3.11a1 3455 (add MAKE_CELL bpo-43693)
+#     Python 3.11a1 3456 (interleave cell args bpo-43693)
 
 #
 # MAGIC must change whenever the bytecode emitted by the compiler may no
@@ -367,7 +368,7 @@ def _write_atomic(path, data, mode=0o666):
 # Whenever MAGIC_NUMBER is changed, the ranges in the magic_values array
 # in PC/launcher.c must also be updated.
 
-MAGIC_NUMBER = (3455).to_bytes(2, 'little') + b'\r\n'
+MAGIC_NUMBER = (3456).to_bytes(2, 'little') + b'\r\n'
 _RAW_MAGIC_NUMBER = int.from_bytes(MAGIC_NUMBER, 'little')  # For import.c
 
 _PYCACHE = '__pycache__'

diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py
@@ -427,9 +427,9 @@ def foo(x):
     return foo
 
 dis_nested_0 = """\
-              0 MAKE_CELL                2 (y)
+              0 MAKE_CELL                0 (y)
 
-%3d           2 LOAD_CLOSURE             2 (y)
+%3d           2 LOAD_CLOSURE             0 (y)
               4 BUILD_TUPLE              1
               6 LOAD_CONST               1 (<code object foo at 0x..., file "%s", line %d>)
               8 LOAD_CONST               2 ('_h.<locals>.foo')
@@ -446,14 +446,14 @@ def foo(x):
 
 dis_nested_1 = """%s
 Disassembly of <code object foo at 0x..., file "%s", line %d>:
-              0 MAKE_CELL                1 (x)
+              0 MAKE_CELL                0 (x)
 
-%3d           2 LOAD_CLOSURE             1 (x)
+%3d           2 LOAD_CLOSURE             0 (x)
               4 BUILD_TUPLE              1
               6 LOAD_CONST               1 (<code object <listcomp> at 0x..., file "%s", line %d>)
               8 LOAD_CONST               2 ('_h.<locals>.foo.<locals>.<listcomp>')
              10 MAKE_FUNCTION            8 (closure)
-             12 LOAD_DEREF               2 (y)
+             12 LOAD_DEREF               1 (y)
              14 GET_ITER
              16 CALL_FUNCTION            1
              18 RETURN_VALUE
@@ -966,19 +966,19 @@ def jumpy():
 
 Instruction = dis.Instruction
 expected_opinfo_outer = [
-  Instruction(opname='MAKE_CELL', opcode=135, arg=3, argval='a', argrepr='a', offset=0, starts_line=None, is_jump_target=False),
-  Instruction(opname='MAKE_CELL', opcode=135, arg=4, argval='b', argrepr='b', offset=2, starts_line=None, is_jump_target=False),
+  Instruction(opname='MAKE_CELL', opcode=135, arg=0, argval='a', argrepr='a', offset=0, starts_line=None, is_jump_target=False),
+  Instruction(opname='MAKE_CELL', opcode=135, arg=1, argval='b', argrepr='b', offset=2, starts_line=None, is_jump_target=False),
   Instruction(opname='LOAD_CONST', opcode=100, arg=8, argval=(3, 4), argrepr='(3, 4)', offset=4, starts_line=2, is_jump_target=False),
-  Instruction(opname='LOAD_CLOSURE', opcode=136, arg=3, argval='a', argrepr='a', offset=6, starts_line=None, is_jump_target=False),
-  Instruction(opname='LOAD_CLOSURE', opcode=136, arg=4, argval='b', argrepr='b', offset=8, starts_line=None, is_jump_target=False),
+  Instruction(opname='LOAD_CLOSURE', opcode=136, arg=0, argval='a', argrepr='a', offset=6, starts_line=None, is_jump_target=False),
+  Instruction(opname='LOAD_CLOSURE', opcode=136, arg=1, argval='b', argrepr='b', offset=8, starts_line=None, is_jump_target=False),
   Instruction(opname='BUILD_TUPLE', opcode=102, arg=2, argval=2, argrepr='', offset=10, starts_line=None, is_jump_target=False),
   Instruction(opname='LOAD_CONST', opcode=100, arg=3, argval=code_object_f, argrepr=repr(code_object_f), offset=12, starts_line=None, is_jump_target=False),
   Instruction(opname='LOAD_CONST', opcode=100, arg=4, argval='outer.<locals>.f', argrepr="'outer.<locals>.f'", offset=14, starts_line=None, is_jump_target=False),
   Instruction(opname='MAKE_FUNCTION', opcode=132, arg=9, argval=9, argrepr='defaults, closure', offset=16, starts_line=None, is_jump_target=False),
   Instruction(opname='STORE_FAST', opcode=125, arg=2, argval='f', argrepr='f', offset=18, starts_line=None, is_jump_target=False),
   Instruction(opname='LOAD_GLOBAL', opcode=116, arg=0, argval='print', argrepr='print', offset=20, starts_line=7, is_jump_target=False),
-  Instruction(opname='LOAD_DEREF', opcode=137, arg=3, argval='a', argrepr='a', offset=22, starts_line=None, is_jump_target=False),
-  Instruction(opname='LOAD_DEREF', opcode=137, arg=4, argval='b', argrepr='b', offset=24, starts_line=None, is_jump_target=False),
+  Instruction(opname='LOAD_DEREF', opcode=137, arg=0, argval='a', argrepr='a', offset=22, starts_line=None, is_jump_target=False),
+  Instruction(opname='LOAD_DEREF', opcode=137, arg=1, argval='b', argrepr='b', offset=24, starts_line=None, is_jump_target=False),
   Instruction(opname='LOAD_CONST', opcode=100, arg=5, argval='', argrepr="''", offset=26, starts_line=None, is_jump_target=False),
   Instruction(opname='LOAD_CONST', opcode=100, arg=6, argval=1, argrepr='1', offset=28, starts_line=None, is_jump_target=False),
   Instruction(opname='BUILD_LIST', opcode=103, arg=0, argval=0, argrepr='', offset=30, starts_line=None, is_jump_target=False),
@@ -991,23 +991,23 @@ def jumpy():
 ]
 
 expected_opinfo_f = [
-  Instruction(opname='MAKE_CELL', opcode=135, arg=3, argval='c', argrepr='c', offset=0, starts_line=None, is_jump_target=False),
-  Instruction(opname='MAKE_CELL', opcode=135, arg=4, argval='d', argrepr='d', offset=2, starts_line=None, is_jump_target=False),
+  Instruction(opname='MAKE_CELL', opcode=135, arg=0, argval='c', argrepr='c', offset=0, starts_line=None, is_jump_target=False),
+  Instruction(opname='MAKE_CELL', opcode=135, arg=1, argval='d', argrepr='d', offset=2, starts_line=None, is_jump_target=False),
   Instruction(opname='LOAD_CONST', opcode=100, arg=5, argval=(5, 6), argrepr='(5, 6)', offset=4, starts_line=3, is_jump_target=False),
-  Instruction(opname='LOAD_CLOSURE', opcode=136, arg=5, argval='a', argrepr='a', offset=6, starts_line=None, is_jump_target=False),
-  Instruction(opname='LOAD_CLOSURE', opcode=136, arg=6, argval='b', argrepr='b', offset=8, starts_line=None, is_jump_target=False),
-  Instruction(opname='LOAD_CLOSURE', opcode=136, arg=3, argval='c', argrepr='c', offset=10, starts_line=None, is_jump_target=False),
-  Instruction(opname='LOAD_CLOSURE', opcode=136, arg=4, argval='d', argrepr='d', offset=12, starts_line=None, is_jump_target=False),
+  Instruction(opname='LOAD_CLOSURE', opcode=136, arg=3, argval='a', argrepr='a', offset=6, starts_line=None, is_jump_target=False),
+  Instruction(opname='LOAD_CLOSURE', opcode=136, arg=4, argval='b', argrepr='b', offset=8, starts_line=None, is_jump_target=False),
+  Instruction(opname='LOAD_CLOSURE', opcode=136, arg=0, argval='c', argrepr='c', offset=10, starts_line=None, is_jump_target=False),
+  Instruction(opname='LOAD_CLOSURE', opcode=136, arg=1, argval='d', argrepr='d', offset=12, starts_line=None, is_jump_target=False),
   Instruction(opname='BUILD_TUPLE', opcode=102, arg=4, argval=4, argrepr='', offset=14, starts_line=None, is_jump_target=False),
   Instruction(opname='LOAD_CONST', opcode=100, arg=3, argval=code_object_inner, argrepr=repr(code_object_inner), offset=16, starts_line=None, is_jump_target=False),
   Instruction(opname='LOAD_CONST', opcode=100, arg=4, argval='outer.<locals>.f.<locals>.inner', argrepr="'outer.<locals>.f.<locals>.inner'", offset=18, starts_line=None, is_jump_target=False),
   Instruction(opname='MAKE_FUNCTION', opcode=132, arg=9, argval=9, argrepr='defaults, closure', offset=20, starts_line=None, is_jump_target=False),
   Instruction(opname='STORE_FAST', opcode=125, arg=2, argval='inner', argrepr='inner', offset=22, starts_line=None, is_jump_target=False),
   Instruction(opname='LOAD_GLOBAL', opcode=116, arg=0, argval='print', argrepr='print', offset=24, starts_line=5, is_jump_target=False),
-  Instruction(opname='LOAD_DEREF', opcode=137, arg=5, argval='a', argrepr='a', offset=26, starts_line=None, is_jump_target=False),
-  Instruction(opname='LOAD_DEREF', opcode=137, arg=6, argval='b', argrepr='b', offset=28, starts_line=None, is_jump_target=False),
-  Instruction(opname='LOAD_DEREF', opcode=137, arg=3, argval='c', argrepr='c', offset=30, starts_line=None, is_jump_target=False),
-  Instruction(opname='LOAD_DEREF', opcode=137, arg=4, argval='d', argrepr='d', offset=32, starts_line=None, is_jump_target=False),
+  Instruction(opname='LOAD_DEREF', opcode=137, arg=3, argval='a', argrepr='a', offset=26, starts_line=None, is_jump_target=False),
+  Instruction(opname='LOAD_DEREF', opcode=137, arg=4, argval='b', argrepr='b', offset=28, starts_line=None, is_jump_target=False),
+  Instruction(opname='LOAD_DEREF', opcode=137, arg=0, argval='c', argrepr='c', offset=30, starts_line=None, is_jump_target=False),
+  Instruction(opname='LOAD_DEREF', opcode=137, arg=1, argval='d', argrepr='d', offset=32, starts_line=None, is_jump_target=False),
   Instruction(opname='CALL_FUNCTION', opcode=131, arg=4, argval=4, argrepr='', offset=34, starts_line=None, is_jump_target=False),
   Instruction(opname='POP_TOP', opcode=1, arg=None, argval=None, argrepr='', offset=36, starts_line=None, is_jump_target=False),
   Instruction(opname='LOAD_FAST', opcode=124, arg=2, argval='inner', argrepr='inner', offset=38, starts_line=6, is_jump_target=False),

diff --git a/Misc/NEWS.d/next/Core and Builtins/2021-06-07-15-13-44.bpo-43693.c_zDeY.rst b/Misc/NEWS.d/next/Core and Builtins/2021-06-07-15-13-44.bpo-43693.c_zDeY.rst
@@ -0,0 +1,4 @@
+Computation of the offsets of cell variables is done in the compiler instead
+of at runtime. This reduces the overhead of handling cell and free
+variables, especially in the case where a variable is both an argument and
+cell variable.
diff --git 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)/
-Original file line number
+Diff line change
@@ Expand Up / @@ -32,8 +32,6 @@ _PyFrame_GetBuiltins(PyFrameObject *f) @@
     int _PyFrame_TakeLocals(PyFrameObject *f);
-    PyAPI_FUNC(int) _PyFrame_OpAlreadyRan(PyFrameObject *f, int opcode, int oparg);
     #ifdef __cplusplus
     }
     #endif
@@ Expand Down @@