Do not create cells in PyFrame_LocalsToFast().

Author: ericsnowcurrently

Include/internal/pycore_frame.h

@@ -32,6 +32,8 @@ _PyFrame_GetBuiltins(PyFrameObject *f)
 
 int _PyFrame_TakeLocals(PyFrameObject *f);
 
+PyAPI_FUNC(int) _PyFrame_OpAlreadyRan(PyFrameObject *f, int opcode, int oparg);
+
 #ifdef __cplusplus
 }
 #endif

Lib/test/test_scope.py

@@ -176,6 +176,57 @@ def bar():
         self.assertEqual(foo(a=42), 50)
         self.assertEqual(foo(), 25)
 
+    def testCellIsArgAndEscapes(self):
+        # We need to be sure that a cell passed in as an arg still
+        # gets wrapped in a new cell if the arg escapes into an
+        # inner function (closure).
+
+        def external():
+            value = 42
+            def inner():
+                return value
+            cell, = inner.__closure__
+            return cell
+        cell_ext = external()
+
+        def spam(arg):
+            def eggs():
+                return arg
+            return eggs
+
+        eggs = spam(cell_ext)
+        cell_closure, = eggs.__closure__
+        cell_eggs = eggs()
+
+        self.assertIs(cell_eggs, cell_ext)
+        self.assertIsNot(cell_eggs, cell_closure)
+
+    def testCellIsLocalAndEscapes(self):
+        # We need to be sure that a cell bound to a local still
+        # gets wrapped in a new cell if the local escapes into an
+        # inner function (closure).
+
+        def external():
+            value = 42
+            def inner():
+                return value
+            cell, = inner.__closure__
+            return cell
+        cell_ext = external()
+
+        def spam(arg):
+            cell = arg
+            def eggs():
+                return cell
+            return eggs
+
+        eggs = spam(cell_ext)
+        cell_closure, = eggs.__closure__
+        cell_eggs = eggs()
+
+        self.assertIs(cell_eggs, cell_ext)
+        self.assertIsNot(cell_eggs, cell_closure)
+
     def testRecursion(self):
 
         def f(x):

Objects/frameobject.c

@@ -918,6 +918,19 @@ PyFrame_New(PyThreadState *tstate, PyCodeObject *code,
     return f;
 }
 
+int
+_PyFrame_OpAlreadyRan(PyFrameObject *f, int opcode, int oparg)
+{
+    const _Py_CODEUNIT *code =
+        (const _Py_CODEUNIT *)PyBytes_AS_STRING(f->f_code->co_code);
+    for (int i = 0; i < f->f_lasti; i++) {
+        if (_Py_OPCODE(code[i]) == opcode && _Py_OPARG(code[i]) == oparg) {
+            return 1;
+        }
+    }
+    return 0;
+}
+
 int
 PyFrame_FastToLocalsWithError(PyFrameObject *f)
 {
@@ -966,32 +979,40 @@ PyFrame_FastToLocalsWithError(PyFrameObject *f)
 
         PyObject *name = PyTuple_GET_ITEM(co->co_localsplusnames, i);
         PyObject *value = fast[i];
-        if (kind & (CO_FAST_CELL | CO_FAST_FREE)) {
+        int cellargoffset = CO_CELL_NOT_AN_ARG;
+        if (co->co_cell2arg != NULL) {
+            cellargoffset = co->co_cell2arg[i - co->co_nlocals];
+        }
+        if (kind & CO_FAST_FREE) {
+            // The cell was set by _PyEval_MakeFrameVector() from
+            // the function's closure.
+            assert(value != NULL && PyCell_Check(value));
+            value = PyCell_GET(value);
+        }
+        else if (kind & CO_FAST_CELL) {
+            // Note that no *_DEREF ops can happen before MAKE_CELL
+            // executes.  So there's no need to duplicate the work
+            // that MAKE_CELL would otherwise do later, if it hasn't
+            // run yet.
             if (value != NULL) {
-                // MAKE_CELL must have executed already.
-                assert(PyCell_Check(value));
-                value = PyCell_GET(value);
+                if (PyCell_Check(value) &&
+                        _PyFrame_OpAlreadyRan(f, MAKE_CELL, i)) {
+                    // (likely) MAKE_CELL must have executed already.
+                    value = PyCell_GET(value);
+                }
+                // (unlikely) Otherwise it must be an initial value set
+                // by an earlier call to PyFrame_FastToLocals().
             }
             else {
-                // MAKE_CELL hasn't executed yet.  (This is unlikely.)
-                PyObject *cell = PyCell_New(NULL);
-                if (cell == NULL) {
-                    PyErr_Clear();
-                }
-                else {
-                    if (co->co_cell2arg != NULL) {
-                        int argoffset = co->co_cell2arg[i - co->co_nlocals];
-                        if (argoffset != CO_CELL_NOT_AN_ARG) {
-                            // It will have been set in initialize_locals().
-                            value = fast[argoffset];
-                            assert(value != NULL);
-                            Py_INCREF(value);
-                            PyCell_SET(cell, value);
-                            // Clear the arg slot.
-                            Py_SETREF(fast[argoffset], NULL);
-                        }
-                    }
-                    fast[i] = cell;
+                // (unlikely) MAKE_CELL hasn't executed yet.
+                if (cellargoffset != CO_CELL_NOT_AN_ARG) {
+                    // It is an arg that escapes into an inner
+                    // function so we use the initial value that
+                    // was already set by _PyEval_MakeFrameVector().
+                    // Normally the arg value would always be set.
+                    // However, it can be NULL if it was deleted via
+                    // PyFrame_LocalsToFast().
+                    value = fast[cellargoffset];
                 }
             }
         }
@@ -1064,36 +1085,66 @@ PyFrame_LocalsToFast(PyFrameObject *f, int clear)
             }
         }
         PyObject *oldvalue = fast[i];
-        if (kind & (CO_FAST_CELL | CO_FAST_FREE)) {
-            if (oldvalue != NULL) {
-                // MAKE_CELL must have executed already.
+        int cellargoffset = CO_CELL_NOT_AN_ARG;
+        if (co->co_cell2arg != NULL) {
+            cellargoffset = co->co_cell2arg[i - co->co_nlocals];
+        }
+        PyObject *cell = NULL;
+        if (kind == CO_FAST_FREE) {
+            // The cell was cell by _PyEval_MakeFrameVector() from
+            // the function's closure.
+            assert(oldvalue != NULL && PyCell_Check(oldvalue));
+            cell = oldvalue;
+        }
+        else if (kind & CO_FAST_CELL && oldvalue != NULL) {
+            if (cellargoffset != CO_CELL_NOT_AN_ARG) {
+                // (likely) MAKE_CELL must have executed already.
+                // It's the cell for an arg.
                 assert(PyCell_Check(oldvalue));
-                Py_XDECREF(PyCell_GET(oldvalue));
-                Py_XINCREF(value);
-                PyCell_SET(oldvalue, value);
+                cell = oldvalue;
             }
             else {
-                // MAKE_CELL hasn't executed yet.  (This is unlikely.)
-                PyObject *cell = PyCell_New(value);
-                if (cell == NULL) {
-                    PyErr_Clear();
+                if (PyCell_Check(oldvalue) &&
+                        _PyFrame_OpAlreadyRan(f, MAKE_CELL, i)) {
+                    // (likely) MAKE_CELL must have executed already.
+                    cell = oldvalue;
                 }
-                else {
-                    fast[i] = cell;
-                    // Clear the corresponding arg if there is one.
-                    if (co->co_cell2arg != NULL) {
-                        int argoffset = co->co_cell2arg[i - co->co_nlocals];
-                        if (argoffset != CO_CELL_NOT_AN_ARG) {
-                            // It will have been set in initialize_locals().
-                            assert(fast[argoffset] != NULL);
-                            Py_SETREF(fast[argoffset], NULL);
-                        }
-                    }
+                // (unlikely) Otherwise, it must have been set to some
+                // initial value by an earlier call to PyFrame_LocalsToFast().
+            }
+        }
+        if (cell != NULL) {
+            PyObject *oldvalue = PyCell_GET(cell);
+            if (value != oldvalue) {
+                Py_XDECREF(oldvalue);
+                Py_XINCREF(value);
+                PyCell_SET(cell, value);
+            }
+        }
+        else {
+            int offset = i;
+            if (kind & CO_FAST_CELL) {
+                // (unlikely) MAKE_CELL hasn't executed yet.
+                // Note that there is no need to create the cell that
+                // MAKE_CELL would otherwise create later, since no
+                // *_DEREF ops can happen before MAKE_CELL has run.
+                if (cellargoffset != CO_CELL_NOT_AN_ARG) {
+                    // It's the cell for an arg.
+                    // Replace the initial value that was set by
+                    // _PyEval_MakeFrameVector().
+                    // Normally the arg value would always be set.
+                    // However, it can be NULL if it was deleted
+                    // via an earlier PyFrame_LocalsToFast() call.
+                    offset = cellargoffset;
+                    oldvalue = fast[offset];
                 }
+                // Otherwise set an initial value for MAKE_CELL to use
+                // when it runs later.
+            }
+            if (value != oldvalue) {
+                Py_XINCREF(value);
+                Py_XSETREF(fast[offset], value);
             }
-        } else if (fast[i] != value) {
-            Py_XINCREF(value);
-            Py_XSETREF(oldvalue, value);
         }
         Py_XDECREF(value);
     }

Objects/typeobject.c

@@ -11,6 +11,8 @@
 #include "pycore_pystate.h"       // _PyThreadState_GET()
 #include "pycore_unionobject.h"   // _Py_Union(), _Py_union_type_or
 #include "frameobject.h"
+#include "pycore_frame.h"         // _PyFrame_OpAlreadyRan
+#include "opcode.h"               // MAKE_CELL
 #include "structmember.h"         // PyMemberDef
 
 #include <ctype.h>
@@ -8883,8 +8885,10 @@ super_init_without_args(PyFrameObject *f, PyCodeObject *co,
         for (i = 0; i < co->co_ncellvars; i++) {
             if (co->co_cell2arg[i] == 0) {
                 PyObject *cell = f->f_localsptr[co->co_nlocals + i];
-                assert(PyCell_Check(cell));
-                obj = PyCell_GET(cell);
+                if (PyCell_Check(cell) &&
+                        _PyFrame_OpAlreadyRan(f, MAKE_CELL, 0)) {
+                    obj = PyCell_GET(cell);
+                }
                 break;
             }
         }

Python/ceval.c

@@ -3077,27 +3077,27 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, PyFrameObject *f, int throwflag)
         }
 
         case TARGET(MAKE_CELL): {
-            if (GETLOCAL(oparg) != NULL) {
-                // A cell must have been set in PyFrame_LocalsToFast().
-                // (This is unlikely to happen.)
-                assert(PyCell_Check(GETLOCAL(oparg)));
-                DISPATCH();
-            }
-            PyObject *cell = PyCell_New(NULL);
+            PyObject *initial = GETLOCAL(oparg);
+            // Normally initial would be NULL.  However, it
+            // might have been set to an initial value during
+            // a call to PyFrame_LocalsToFast().
+            PyObject *cell = PyCell_New(initial);
             if (cell == NULL) {
                 goto error;
             }
             /* If it is an arg then copy the arg into the cell. */
-            if (co->co_cell2arg != NULL) {
+            if (initial == NULL && co->co_cell2arg != NULL) {
                 int argoffset = co->co_cell2arg[oparg - co->co_nlocals];
                 if (argoffset != CO_CELL_NOT_AN_ARG) {
                     PyObject *arg = GETLOCAL(argoffset);
-                    // It will have been set in initialize_locals().
-                    assert(arg != NULL);
-                    Py_INCREF(arg);
-                    PyCell_SET(cell, arg);
-                    /* Clear the local copy. */
-                    SETLOCAL(argoffset, NULL);
+                    // It will have been set in initialize_locals() but
+                    // may have been deleted PyFrame_LocalsToFast().
+                    if (arg != NULL) {;
+                        Py_INCREF(arg);
+                        PyCell_SET(cell, arg);
+                        /* Clear the local copy. */
+                        SETLOCAL(argoffset, NULL);
+                    }
                 }
             }
             SETLOCAL(oparg, cell);