map cells to arg slots at code creation time (closes #12399)

This removes nested loops in PyEval_EvalCodeEx.

Include/code.h

@@ -22,6 +22,7 @@ typedef struct {
     PyObject *co_freevars;	/* tuple of strings (free variable names) */
     PyObject *co_cellvars;      /* tuple of strings (cell variable names) */
     /* The rest doesn't count for hash or comparisons */
+    unsigned char *co_cell2arg; /* Maps cell vars which are arguments. */
     PyObject *co_filename;	/* unicode (where it was loaded from) */
     PyObject *co_name;		/* unicode (name, for reference) */
     int co_firstlineno;		/* first source line number */
@@ -57,6 +58,11 @@ typedef struct {
 
 #define CO_FUTURE_BARRY_AS_BDFL  0x40000
 
+/* This value is found in the co_cell2arg array when the associated cell
+   variable does not correspond to an argument. The maximum number of
+   arguments is 255 (indexed up to 254), so 255 work as a special flag.*/
+#define CO_CELL_NOT_AN_ARG 255
+
 /* This should be defined if a future statement modifies the syntax.
    For example, when a keyword is added.
 */

Lib/test/test_sys.py

@@ -665,7 +665,7 @@ class SizeofTest(unittest.TestCase):
             return inner
         check(get_cell().__closure__[0], size(h + 'P'))
         # code
-        check(get_cell().__code__, size(h + '5i8Pi3P'))
+        check(get_cell().__code__, size(h + '5i9Pi3P'))
         # complex
         check(complex(0,1), size(h + '2d'))
         # method_descriptor (descriptor object)

Objects/codeobject.c

@@ -51,7 +51,8 @@ PyCode_New(int argcount, int kwonlyargcount,
            PyObject *lnotab)
 {
     PyCodeObject *co;
-    Py_ssize_t i;
+    unsigned char *cell2arg = NULL;
+    Py_ssize_t i, n_cellvars;
 
     /* Check argument types */
     if (argcount < 0 || kwonlyargcount < 0 || nlocals < 0 ||
@@ -68,12 +69,13 @@ PyCode_New(int argcount, int kwonlyargcount,
         PyErr_BadInternalCall();
         return NULL;
     }
+    n_cellvars = PyTuple_GET_SIZE(cellvars);
     intern_strings(names);
     intern_strings(varnames);
     intern_strings(freevars);
     intern_strings(cellvars);
     /* Intern selected string constants */
-    for (i = PyTuple_Size(consts); --i >= 0; ) {
+    for (i = PyTuple_GET_SIZE(consts); --i >= 0; ) {
         PyObject *v = PyTuple_GetItem(consts, i);
         if (!PyUnicode_Check(v))
             continue;
@@ -81,8 +83,40 @@ PyCode_New(int argcount, int kwonlyargcount,
             continue;
         PyUnicode_InternInPlace(&PyTuple_GET_ITEM(consts, i));
     }
+    /* Create mapping between cells and arguments if needed. */
+    if (n_cellvars) {
+        Py_ssize_t total_args = argcount + kwonlyargcount +
+            ((flags & CO_VARARGS) != 0) + ((flags & CO_VARKEYWORDS) != 0);
+        Py_ssize_t alloc_size = sizeof(unsigned char) * n_cellvars;
+        int used_cell2arg = 0;
+        cell2arg = PyMem_MALLOC(alloc_size);
+        if (cell2arg == NULL)
+            return NULL;
+        memset(cell2arg, CO_CELL_NOT_AN_ARG, alloc_size);
+        /* Find cells which are also arguments. */
+        for (i = 0; i < n_cellvars; i++) {
+            Py_ssize_t j;
+            PyObject *cell = PyTuple_GET_ITEM(cellvars, i);
+            for (j = 0; j < total_args; j++) {
+                PyObject *arg = PyTuple_GET_ITEM(varnames, j);
+                if (!PyUnicode_Compare(cell, arg)) {
+                    cell2arg[i] = j;
+                    used_cell2arg = 1;
+                    break;
+                }
+            }
+        }
+        if (!used_cell2arg) {
+            PyMem_FREE(cell2arg);
+            cell2arg = NULL;
+        }
+    }
     co = PyObject_NEW(PyCodeObject, &PyCode_Type);
-    if (co != NULL) {
+    if (co == NULL) {
+        if (cell2arg)
+            PyMem_FREE(cell2arg);
+        return NULL;
+    }
     co->co_argcount = argcount;
     co->co_kwonlyargcount = kwonlyargcount;
     co->co_nlocals = nlocals;
@@ -100,6 +134,7 @@ PyCode_New(int argcount, int kwonlyargcount,
     co->co_freevars = freevars;
     Py_INCREF(cellvars);
     co->co_cellvars = cellvars;
+    co->co_cell2arg = cell2arg;
     Py_INCREF(filename);
     co->co_filename = filename;
     Py_INCREF(name);
@@ -109,7 +144,6 @@ PyCode_New(int argcount, int kwonlyargcount,
     co->co_lnotab = lnotab;
     co->co_zombieframe = NULL;
     co->co_weakreflist = NULL;
-    }
     return co;
 }
 
@@ -330,6 +364,8 @@ code_dealloc(PyCodeObject *co)
     Py_XDECREF(co->co_filename);
     Py_XDECREF(co->co_name);
     Py_XDECREF(co->co_lnotab);
+    if (co->co_cell2arg != NULL)
+        PyMem_FREE(co->co_cell2arg);
     if (co->co_zombieframe != NULL)
         PyObject_GC_Del(co->co_zombieframe);
     if (co->co_weakreflist != NULL)

Python/ceval.c

@@ -3357,57 +3357,25 @@ PyEval_EvalCodeEx(PyObject *_co, PyObject *globals, PyObject *locals,
     }
 
     /* Allocate and initialize storage for cell vars, and copy free
-       vars into frame.  This isn't too efficient right now. */
-    if (PyTuple_GET_SIZE(co->co_cellvars)) {
-        int i, j, nargs, found;
-        Py_UNICODE *cellname, *argname;
-        PyObject *c;
-
-        nargs = total_args;
-        if (co->co_flags & CO_VARARGS)
-            nargs++;
-        if (co->co_flags & CO_VARKEYWORDS)
-            nargs++;
-
-        /* Initialize each cell var, taking into account
-           cell vars that are initialized from arguments.
-
-           Should arrange for the compiler to put cellvars
-           that are arguments at the beginning of the cellvars
-           list so that we can march over it more efficiently?
-        */
+       vars into frame. */
     for (i = 0; i < PyTuple_GET_SIZE(co->co_cellvars); ++i) {
-            cellname = PyUnicode_AS_UNICODE(
-                PyTuple_GET_ITEM(co->co_cellvars, i));
-            found = 0;
-            for (j = 0; j < nargs; j++) {
-                argname = PyUnicode_AS_UNICODE(
-                    PyTuple_GET_ITEM(co->co_varnames, j));
-                if (Py_UNICODE_strcmp(cellname, argname) == 0) {
-                    c = PyCell_New(GETLOCAL(j));
-                    if (c == NULL)
-                        goto fail;
-                    GETLOCAL(co->co_nlocals + i) = c;
-                    found = 1;
-                    break;
-                }
-            }
-            if (found == 0) {
+        PyObject *c;
+        int arg;
+        /* Possibly account for the cell variable being an argument. */
+        if (co->co_cell2arg != NULL &&
+            (arg = co->co_cell2arg[i]) != CO_CELL_NOT_AN_ARG)
+            c = PyCell_New(GETLOCAL(arg));
+        else
             c = PyCell_New(NULL);
         if (c == NULL)
             goto fail;
         SETLOCAL(co->co_nlocals + i, c);
     }
-        }
-    }
-    if (PyTuple_GET_SIZE(co->co_freevars)) {
-        int i;
     for (i = 0; i < PyTuple_GET_SIZE(co->co_freevars); ++i) {
         PyObject *o = PyTuple_GET_ITEM(closure, i);
         Py_INCREF(o);
         freevars[PyTuple_GET_SIZE(co->co_cellvars) + i] = o;
     }
-    }
 
     if (co->co_flags & CO_GENERATOR) {
         /* Don't need to keep the reference to f_back, it will be set