Add co_fastlocalnames and co_fastlocalkinds.

ericsnowcurrently · Jun 3, 2021 · May 20, 2021 · May 26, 2021 · May 21, 2021 · May 21, 2021
commit a5edf53b8520708f7a517bbc69ad1731ea881329
diff --git a/Include/cpython/code.h b/Include/cpython/code.h
@@ -3,6 +3,8 @@
 #endif
 
 typedef uint16_t _Py_CODEUNIT;
+// Each oparg must fit in the second half of _Py_CODEUNIT, hence 8 bits.
+#define _Py_MAX_OPARG 255
 
 #ifdef WORDS_BIGENDIAN
 #  define _Py_OPCODE(word) ((word) >> 8)
@@ -14,6 +16,11 @@ typedef uint16_t _Py_CODEUNIT;
 
 typedef struct _PyOpcache _PyOpcache;
 
+
+// These are duplicated from pycore_code.h.
+typedef unsigned char _PyFastLocalKind;
+typedef _PyFastLocalKind _PyFastLocalKinds[_Py_MAX_OPARG + 1];
+
 /* Bytecode object */
 struct PyCodeObject {
     PyObject_HEAD
@@ -53,9 +60,8 @@ struct PyCodeObject {
     int co_kwonlyargcount;      /* #keyword only arguments */
     int co_stacksize;           /* #entries needed for evaluation stack */
     int co_firstlineno;         /* first source line number */
-    PyObject *co_varnames;      /* tuple of strings (local variable names) */
-    PyObject *co_cellvars;      /* tuple of strings (cell variable names) */
-    PyObject *co_freevars;      /* tuple of strings (free variable names) */
+    PyObject *co_fastlocalnames;  /* tuple mapping fastlocalsplus to names */
+    _PyFastLocalKinds co_fastlocalkinds; /* array mapping to local kinds */
     PyObject *co_filename;      /* unicode (where it was loaded from) */
     PyObject *co_name;          /* unicode (name, for reference) */
     PyObject *co_linetable;     /* string (encoding addr<->lineno mapping) See
@@ -70,6 +76,10 @@ struct PyCodeObject {
     int co_nlocals;             /* number of local variables */
     int co_ncellvars;           /* number of cell variables */
     int co_nfreevars;           /* number of free variables */
+    // cached values
+    PyObject *co_varnames;      /* tuple of strings (local variable names) */
+    PyObject *co_cellvars;      /* tuple of strings (cell variable names) */
+    PyObject *co_freevars;      /* tuple of strings (free variable names) */
 
     /* The remaining fields are zeroed out on new code objects. */
 

diff --git a/Include/internal/pycore_code.h b/Include/internal/pycore_code.h
@@ -26,6 +26,25 @@ struct _PyOpcache {
 };
 
 
+// We would use an enum if C let us specify the storage type.
+typedef unsigned char _PyFastLocalKind;
+// Note that these all fit within _PyFastLocalKind, as do combinations.
+#define CO_FAST_POSONLY     0x01
+#define CO_FAST_POSORKW     0x02
+#define CO_FAST_VARARGS     0x04
+#define CO_FAST_KWONLY      0x08
+#define CO_FAST_VARKWARGS   0x10
+#define CO_FAST_LOCALONLY   0x20
+#define CO_FAST_CELL        0x40
+#define CO_FAST_FREE        0x80
+
+#define CO_FAST_ARG (CO_FAST_POSONLY | CO_FAST_POSORKW | CO_FAST_VARARGS | \
+                     CO_FAST_KWONLY | CO_FAST_VARKWARGS)
+#define CO_FAST_LOCAL (CO_FAST_ARG | CO_FAST_LOCALONLY)
+#define CO_FAST_ANY (CO_FAST_LOCAL | CO_FAST_CELL | CO_FAST_FREE)
+
+typedef _PyFastLocalKind _PyFastLocalKinds[_Py_MAX_OPARG + 1];
+
 struct _PyCodeConstructor {
     /* metadata */
     PyObject *filename;
@@ -42,9 +61,8 @@ struct _PyCodeConstructor {
     PyObject *names;
 
     /* mapping frame offsets to information */
-    PyObject *varnames;
-    PyObject *cellvars;
-    PyObject *freevars;
+    PyObject *fastlocalnames;
+    _PyFastLocalKinds fastlocalkinds;
 
     /* args (within varnames) */
     int argcount;
@@ -56,6 +74,12 @@ struct _PyCodeConstructor {
 
     /* used by the eval loop */
     PyObject *exceptiontable;
+
+    // Everything else could be computed later from fastlocal*.
+    // XXX Drop these.
+    PyObject *varnames;
+    PyObject *cellvars;
+    PyObject *freevars;
 };
 
 // Using an "arguments struct" like this is helpful for maintainability
@@ -75,7 +99,6 @@ PyAPI_FUNC(PyCodeObject *) _PyCode_New(struct _PyCodeConstructor *);
 
 int _PyCode_InitOpcache(PyCodeObject *co);
 
-
 #ifdef __cplusplus
 }
 #endif

@@ -164,6 +164,7 @@ _PyCode_Validate(struct _PyCodeConstructor *con)
         con->code == NULL || !PyBytes_Check(con->code) ||
         con->consts == NULL || !PyTuple_Check(con->consts) ||
         con->names == NULL || !PyTuple_Check(con->names) ||
+        (con->fastlocalnames != NULL && !PyTuple_Check(con->fastlocalnames)) ||
         con->varnames == NULL || !PyTuple_Check(con->varnames) ||
         con->freevars == NULL || !PyTuple_Check(con->freevars) ||
         con->cellvars == NULL || !PyTuple_Check(con->cellvars) ||
@@ -203,6 +204,11 @@ _PyCode_Validate(struct _PyCodeConstructor *con)
 static void
 init_code(PyCodeObject *co, struct _PyCodeConstructor *con)
 {
+    int nlocals = (int)PyTuple_GET_SIZE(con->varnames);
+    int ncellvars = (int)PyTuple_GET_SIZE(con->cellvars);
+    int nfreevars = (int)PyTuple_GET_SIZE(con->freevars);
+    int nlocalsplus = nlocals + ncellvars + nfreevars;
+
     Py_INCREF(con->filename);
     co->co_filename = con->filename;
     Py_INCREF(con->name);
@@ -220,12 +226,9 @@ init_code(PyCodeObject *co, struct _PyCodeConstructor *con)
     Py_INCREF(con->names);
     co->co_names = con->names;
 
-    Py_INCREF(con->varnames);
-    co->co_varnames = con->varnames;
-    Py_INCREF(con->cellvars);
-    co->co_cellvars = con->cellvars;
-    Py_INCREF(con->freevars);
-    co->co_freevars = con->freevars;
+    Py_INCREF(con->fastlocalnames);
+    co->co_fastlocalnames = con->fastlocalnames;
+    memcpy(co->co_fastlocalkinds, con->fastlocalkinds, nlocalsplus);
 
     co->co_argcount = con->argcount;
     co->co_posonlyargcount = con->posonlyargcount;
@@ -238,10 +241,16 @@ init_code(PyCodeObject *co, struct _PyCodeConstructor *con)
 
     /* derived values */
     co->co_cell2arg = NULL;  // This will be set soon.
-    co->co_nlocals = (int)PyTuple_GET_SIZE(con->varnames);
-    co->co_ncellvars = (int)PyTuple_GET_SIZE(con->cellvars);
-    co->co_nfreevars = (int)PyTuple_GET_SIZE(con->freevars);
-    co->co_nlocalsplus = co->co_nlocals + co->co_ncellvars + co->co_nfreevars;
+    co->co_nlocalsplus = nlocalsplus;
+    co->co_nlocals = nlocals;
+    co->co_ncellvars = ncellvars;
+    co->co_nfreevars = nfreevars;
+    Py_INCREF(con->varnames);
+    co->co_varnames = con->varnames;
+    Py_INCREF(con->cellvars);
+    co->co_cellvars = con->cellvars;
+    Py_INCREF(con->freevars);
+    co->co_freevars = con->freevars;
 
     /* not set */
     co->co_weakreflist = NULL;
@@ -271,6 +280,9 @@ _PyCode_New(struct _PyCodeConstructor *con)
     if (intern_string_constants(con->consts, NULL) < 0) {
         return NULL;
     }
+    if (intern_strings(con->fastlocalnames) < 0) {
+        return NULL;
+    }
     if (intern_strings(con->varnames) < 0) {
         return NULL;
     }
@@ -341,6 +353,35 @@ _PyCode_New(struct _PyCodeConstructor *con)
  * the legacy "constructors"
  ******************/
 
+// This is also used in compile.c.
+void
+set_fastlocal_info(Py_ssize_t offset, PyObject *name, _PyFastLocalKind kind,
+                   PyObject *names, _PyFastLocalKinds kinds)
+{
+    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++) {
+            _PyFastLocalKind 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());
+        }
+    }
+}
+
 PyCodeObject *
 PyCode_NewWithPosOnlyArgs(int argcount, int posonlyargcount, int kwonlyargcount,
                           int nlocals, int stacksize, int flags,
@@ -361,9 +402,7 @@ PyCode_NewWithPosOnlyArgs(int argcount, int posonlyargcount, int kwonlyargcount,
         .consts = consts,
         .names = names,
 
-        .varnames = varnames,
-        .cellvars = cellvars,
-        .freevars = freevars,
+        // fastlocalnames and fastlocalkinds are set below.
 
         .argcount = argcount,
         .posonlyargcount = posonlyargcount,
@@ -372,6 +411,10 @@ PyCode_NewWithPosOnlyArgs(int argcount, int posonlyargcount, int kwonlyargcount,
         .stacksize = stacksize,
 
         .exceptiontable = exceptiontable,
+
+        .varnames = varnames,
+        .cellvars = cellvars,
+        .freevars = freevars,
     };
     if (_PyCode_Validate(&con) < 0) {
         return NULL;
@@ -383,7 +426,37 @@ PyCode_NewWithPosOnlyArgs(int argcount, int posonlyargcount, int kwonlyargcount,
         return NULL;
     }
 
-    return _PyCode_New(&con);
+    // Set the "fast locals plus" info.
+    int ncellvars = (int)PyTuple_GET_SIZE(cellvars);
+    int nfreevars = (int)PyTuple_GET_SIZE(freevars);
+    int nlocalsplus = nlocals + ncellvars + nfreevars;
+    con.fastlocalnames = PyTuple_New(nlocalsplus);
+    if (con.fastlocalnames == NULL) {
+        return NULL;
+    }
+    // con->fastlocalkinds is statically allocated so it's ready to go.
+    int  offset = 0;
+    for (int i = 0; i < nlocals; i++, offset++) {
+        PyObject *name = PyTuple_GET_ITEM(varnames, i);
+        set_fastlocal_info(offset, name, CO_FAST_LOCAL,
+                           con.fastlocalnames, con.fastlocalkinds);
+    }
+    for (int i = 0; i < ncellvars; i++, offset++) {
+        PyObject *name = PyTuple_GET_ITEM(cellvars, i);
+        set_fastlocal_info(offset, name, CO_FAST_CELL,
+                           con.fastlocalnames, con.fastlocalkinds);
+    }
+    for (int i = 0; i < nfreevars; i++, offset++) {
+        PyObject *name = PyTuple_GET_ITEM(freevars, i);
+        set_fastlocal_info(offset, name, CO_FAST_FREE,
+                           con.fastlocalnames, con.fastlocalkinds);
+    }
+
+    PyCodeObject *co = _PyCode_New(&con);
+    if (co == NULL) {
+        Py_DECREF(con.fastlocalnames);
+    }
+    return co;
 }
 
 PyCodeObject *
@@ -434,6 +507,7 @@ PyCode_NewEmpty(const char *filename, const char *funcname, int firstlineno)
         .linetable = emptystring,
         .consts = nulltuple,
         .names = nulltuple,
+        .fastlocalnames = nulltuple,
         .varnames = nulltuple,
         .cellvars = nulltuple,
         .freevars = nulltuple,
@@ -1086,8 +1160,6 @@ code_richcompare(PyObject *self, PyObject *other, int op)
     if (!eq) goto unequal;
     eq = co->co_kwonlyargcount == cp->co_kwonlyargcount;
     if (!eq) goto unequal;
-    eq = co->co_nlocals == cp->co_nlocals;
-    if (!eq) goto unequal;
     eq = co->co_flags == cp->co_flags;
     if (!eq) goto unequal;
     eq = co->co_firstlineno == cp->co_firstlineno;
@@ -1111,11 +1183,8 @@ code_richcompare(PyObject *self, PyObject *other, int op)
 
     eq = PyObject_RichCompareBool(co->co_names, cp->co_names, Py_EQ);
     if (eq <= 0) goto unequal;
-    eq = PyObject_RichCompareBool(co->co_varnames, cp->co_varnames, Py_EQ);
-    if (eq <= 0) goto unequal;
-    eq = PyObject_RichCompareBool(co->co_freevars, cp->co_freevars, Py_EQ);
-    if (eq <= 0) goto unequal;
-    eq = PyObject_RichCompareBool(co->co_cellvars, cp->co_cellvars, Py_EQ);
+    eq = PyObject_RichCompareBool(co->co_fastlocalnames,
+                                  cp->co_fastlocalnames, Py_EQ);
     if (eq <= 0) goto unequal;
 
     if (op == Py_EQ)
@@ -1140,7 +1209,7 @@ code_richcompare(PyObject *self, PyObject *other, int op)
 static Py_hash_t
 code_hash(PyCodeObject *co)
 {
-    Py_hash_t h, h0, h1, h2, h3, h4, h5, h6;
+    Py_hash_t h, h0, h1, h2, h3, h4;
     h0 = PyObject_Hash(co->co_name);
     if (h0 == -1) return -1;
     h1 = PyObject_Hash(co->co_code);
@@ -1149,15 +1218,11 @@ code_hash(PyCodeObject *co)
     if (h2 == -1) return -1;
     h3 = PyObject_Hash(co->co_names);
     if (h3 == -1) return -1;
-    h4 = PyObject_Hash(co->co_varnames);
+    h4 = PyObject_Hash(co->co_fastlocalnames);
     if (h4 == -1) return -1;
-    h5 = PyObject_Hash(co->co_freevars);
-    if (h5 == -1) return -1;
-    h6 = PyObject_Hash(co->co_cellvars);
-    if (h6 == -1) return -1;
-    h = h0 ^ h1 ^ h2 ^ h3 ^ h4 ^ h5 ^ h6 ^
+    h = h0 ^ h1 ^ h2 ^ h3 ^ h4 ^
         co->co_argcount ^ co->co_posonlyargcount ^ co->co_kwonlyargcount ^
-        co->co_nlocals ^ co->co_flags;
+        co->co_flags;
     if (h == -1) h = -2;
     return h;
 }
@@ -1204,22 +1269,25 @@ static PyObject *
 code_sizeof(PyCodeObject *co, PyObject *Py_UNUSED(args))
 {
     Py_ssize_t res = _PyObject_SIZE(Py_TYPE(co));
-    _PyCodeObjectExtra *co_extra = (_PyCodeObjectExtra*) co->co_extra;
 
-    if (co->co_cell2arg != NULL && co->co_cellvars != NULL) {
-        res += co->co_ncellvars * sizeof(Py_ssize_t);
-    }
+    _PyCodeObjectExtra *co_extra = (_PyCodeObjectExtra*) co->co_extra;
     if (co_extra != NULL) {
         res += sizeof(_PyCodeObjectExtra) +
                (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_opcache != NULL) {
         assert(co->co_opcache_map != NULL);
         // co_opcache_map
         res += PyBytes_GET_SIZE(co->co_code) / sizeof(_Py_CODEUNIT);
         // co_opcache
         res += co->co_opcache_size * sizeof(_PyOpcache);
     }
+
     return PyLong_FromSsize_t(res);
 }
 

@@ -1065,8 +1065,9 @@ PyFrame_FastToLocalsWithError(PyFrameObject *f)
         */
         if (co->co_flags & CO_OPTIMIZED) {
             if (map_to_dict(co->co_freevars, co->co_nfreevars, locals,
-                            fast + co->co_nlocals + co->co_ncellvars, 1) < 0)
+                            fast + co->co_nlocals + co->co_ncellvars, 1) < 0) {
                 return -1;
+            }
         }
     }
     return 0;