bpo-44525: Copy free variables in bytecode to allow calls to inner functions to be specialized (GH-29595)

* Make internal APIs that take PyFrameConstructor take a PyFunctionObject instead.

* Add reference to function to frame, borrow references to builtins and globals.

* Add COPY_FREE_VARS instruction to allow specialization of calls to inner functions.

5 files changed, 71 insertions, 23 deletions

diff --git a/Objects/call.c b/Objects/call.c
index 5e55518..310a2d7 100644
--- a/Objects/call.c
+++ b/Objects/call.c
@@ -383,16 +383,16 @@ _PyFunction_Vectorcall(PyObject *func, PyObject* const* stack,
                        size_t nargsf, PyObject *kwnames)
 {
     assert(PyFunction_Check(func));
-    PyFrameConstructor *f = PyFunction_AS_FRAME_CONSTRUCTOR(func);
+    PyFunctionObject *f = (PyFunctionObject *)func;
     Py_ssize_t nargs = PyVectorcall_NARGS(nargsf);
     assert(nargs >= 0);
     PyThreadState *tstate = _PyThreadState_GET();
     assert(nargs == 0 || stack != NULL);
-    if (((PyCodeObject *)f->fc_code)->co_flags & CO_OPTIMIZED) {
+    if (((PyCodeObject *)f->func_code)->co_flags & CO_OPTIMIZED) {
         return _PyEval_Vector(tstate, f, NULL, stack, nargs, kwnames);
     }
     else {
-        return _PyEval_Vector(tstate, f, f->fc_globals, stack, nargs, kwnames);
+        return _PyEval_Vector(tstate, f, f->func_globals, stack, nargs, kwnames);
     }
 }
 
diff --git a/Objects/frameobject.c b/Objects/frameobject.c
index 09857c7..2a283b3 100644
--- a/Objects/frameobject.c
+++ b/Objects/frameobject.c
@@ -5,6 +5,7 @@
 #include "pycore_moduleobject.h"  // _PyModule_GetDict()
 #include "pycore_object.h"        // _PyObject_GC_UNTRACK()
 #include "pycore_code.h"          // CO_FAST_LOCAL, etc.
+#include "pycore_function.h"      // _PyFunction_FromConstructor()
 
 #include "frameobject.h"          // PyFrameObject
 #include "pycore_frame.h"
@@ -626,8 +627,7 @@ frame_dealloc(PyFrameObject *f)
         /* Don't clear code object until the end */
         co = frame->f_code;
         frame->f_code = NULL;
-        Py_CLEAR(frame->f_globals);
-        Py_CLEAR(frame->f_builtins);
+        Py_CLEAR(frame->f_func);
         Py_CLEAR(frame->f_locals);
         PyObject **locals = _PyFrame_GetLocalsArray(frame);
         for (int i = 0; i < frame->stacktop; i++) {
@@ -782,16 +782,16 @@ PyTypeObject PyFrame_Type = {
 _Py_IDENTIFIER(__builtins__);
 
 static InterpreterFrame *
-allocate_heap_frame(PyFrameConstructor *con, PyObject *locals)
+allocate_heap_frame(PyFunctionObject *func, PyObject *locals)
 {
-    PyCodeObject *code = (PyCodeObject *)con->fc_code;
+    PyCodeObject *code = (PyCodeObject *)func->func_code;
     int size = code->co_nlocalsplus+code->co_stacksize + FRAME_SPECIALS_SIZE;
     InterpreterFrame *frame = (InterpreterFrame *)PyMem_Malloc(sizeof(PyObject *)*size);
     if (frame == NULL) {
         PyErr_NoMemory();
         return NULL;
     }
-    _PyFrame_InitializeSpecials(frame, con, locals, code->co_nlocalsplus);
+    _PyFrame_InitializeSpecials(frame, func, locals, code->co_nlocalsplus);
     for (Py_ssize_t i = 0; i < code->co_nlocalsplus; i++) {
         frame->localsplus[i] = NULL;
     }
@@ -872,7 +872,12 @@ PyFrame_New(PyThreadState *tstate, PyCodeObject *code,
         .fc_kwdefaults = NULL,
         .fc_closure = NULL
     };
-    InterpreterFrame *frame = allocate_heap_frame(&desc, locals);
+    PyFunctionObject *func = _PyFunction_FromConstructor(&desc);
+    if (func == NULL) {
+        return NULL;
+    }
+    InterpreterFrame *frame = allocate_heap_frame(func, locals);
+    Py_DECREF(func);
     if (frame == NULL) {
         return NULL;
     }
@@ -910,6 +915,18 @@ _PyFrame_FastToLocalsWithError(InterpreterFrame *frame) {
     }
     co = frame->f_code;
     fast = _PyFrame_GetLocalsArray(frame);
+    if (frame->f_lasti < 0 && _Py_OPCODE(co->co_firstinstr[0]) == COPY_FREE_VARS) {
+        /* Free vars have not been initialized -- Do that */
+        PyCodeObject *co = frame->f_code;
+        PyObject *closure = frame->f_func->func_closure;
+        int offset = co->co_nlocals + co->co_nplaincellvars;
+        for (int i = 0; i < co->co_nfreevars; ++i) {
+            PyObject *o = PyTuple_GET_ITEM(closure, i);
+            Py_INCREF(o);
+            frame->localsplus[offset + i] = o;
+        }
+        frame->f_lasti = 0;
+    }
     for (int i = 0; i < co->co_nlocalsplus; i++) {
         _PyLocals_Kind kind = _PyLocals_GetKind(co->co_localspluskinds, i);
 
@@ -929,8 +946,7 @@ _PyFrame_FastToLocalsWithError(InterpreterFrame *frame) {
         PyObject *value = fast[i];
         if (frame->f_state != FRAME_CLEARED) {
             if (kind & CO_FAST_FREE) {
-                // The cell was set when the frame was created from
-                // the function's closure.
+                // The cell was set by COPY_FREE_VARS.
                 assert(value != NULL && PyCell_Check(value));
                 value = PyCell_GET(value);
             }
diff --git a/Objects/funcobject.c b/Objects/funcobject.c
index 5a17038..7891e4f 100644
--- a/Objects/funcobject.c
+++ b/Objects/funcobject.c
@@ -9,6 +9,39 @@
 
 static uint32_t next_func_version = 1;
 
+PyFunctionObject *
+_PyFunction_FromConstructor(PyFrameConstructor *constr)
+{
+
+    PyFunctionObject *op = PyObject_GC_New(PyFunctionObject, &PyFunction_Type);
+    if (op == NULL) {
+        return NULL;
+    }
+    Py_INCREF(constr->fc_globals);
+    op->func_globals = constr->fc_globals;
+    Py_INCREF(constr->fc_builtins);
+    op->func_builtins = constr->fc_builtins;
+    Py_INCREF(constr->fc_name);
+    op->func_name = constr->fc_name;
+    Py_INCREF(constr->fc_qualname);
+    op->func_qualname = constr->fc_qualname;
+    Py_INCREF(constr->fc_code);
+    op->func_code = constr->fc_code;
+    op->func_defaults = NULL;
+    op->func_kwdefaults = NULL;
+    op->func_closure = NULL;
+    Py_INCREF(Py_None);
+    op->func_doc = Py_None;
+    op->func_dict = NULL;
+    op->func_weakreflist = NULL;
+    op->func_module = NULL;
+    op->func_annotations = NULL;
+    op->vectorcall = _PyFunction_Vectorcall;
+    op->func_version = 0;
+    _PyObject_GC_TRACK(op);
+    return op;
+}
+
 PyObject *
 PyFunction_NewWithQualName(PyObject *code, PyObject *globals, PyObject *qualname)
 {
diff --git a/Objects/genobject.c b/Objects/genobject.c
index c899ed6..24d5f35 100644
--- a/Objects/genobject.c
+++ b/Objects/genobject.c
@@ -188,7 +188,6 @@ gen_send_ex2(PyGenObject *gen, PyObject *arg, PyObject **presult,
     }
 
     assert(_PyFrame_IsRunnable(frame));
-    assert(frame->f_lasti >= 0 || ((unsigned char *)PyBytes_AS_STRING(gen->gi_code->co_code))[0] == GEN_START);
     /* Push arg onto the frame's value stack */
     result = arg ? arg : Py_None;
     Py_INCREF(result);
@@ -841,7 +840,7 @@ PyTypeObject PyGen_Type = {
 };
 
 static PyObject *
-make_gen(PyTypeObject *type, PyFrameConstructor *con, InterpreterFrame *frame)
+make_gen(PyTypeObject *type, PyFunctionObject *func, InterpreterFrame *frame)
 {
     PyGenObject *gen = PyObject_GC_New(PyGenObject, type);
     if (gen == NULL) {
@@ -858,13 +857,13 @@ make_gen(PyTypeObject *type, PyFrameConstructor *con, InterpreterFrame *frame)
     gen->gi_exc_state.exc_value = NULL;
     gen->gi_exc_state.exc_traceback = NULL;
     gen->gi_exc_state.previous_item = NULL;
-    if (con->fc_name != NULL)
-        gen->gi_name = con->fc_name;
+    if (func->func_name != NULL)
+        gen->gi_name = func->func_name;
     else
         gen->gi_name = gen->gi_code->co_name;
     Py_INCREF(gen->gi_name);
-    if (con->fc_qualname != NULL)
-        gen->gi_qualname = con->fc_qualname;
+    if (func->func_qualname != NULL)
+        gen->gi_qualname = func->func_qualname;
     else
         gen->gi_qualname = gen->gi_name;
     Py_INCREF(gen->gi_qualname);
@@ -876,17 +875,17 @@ static PyObject *
 compute_cr_origin(int origin_depth);
 
 PyObject *
-_Py_MakeCoro(PyFrameConstructor *con, InterpreterFrame *frame)
+_Py_MakeCoro(PyFunctionObject *func, InterpreterFrame *frame)
 {
-    int coro_flags = ((PyCodeObject *)con->fc_code)->co_flags &
+    int coro_flags = ((PyCodeObject *)func->func_code)->co_flags &
         (CO_GENERATOR | CO_COROUTINE | CO_ASYNC_GENERATOR);
     assert(coro_flags);
     if (coro_flags == CO_GENERATOR) {
-        return make_gen(&PyGen_Type, con, frame);
+        return make_gen(&PyGen_Type, func, frame);
     }
     if (coro_flags == CO_ASYNC_GENERATOR) {
         PyAsyncGenObject *o;
-        o = (PyAsyncGenObject *)make_gen(&PyAsyncGen_Type, con, frame);
+        o = (PyAsyncGenObject *)make_gen(&PyAsyncGen_Type, func, frame);
         if (o == NULL) {
             return NULL;
         }
@@ -897,7 +896,7 @@ _Py_MakeCoro(PyFrameConstructor *con, InterpreterFrame *frame)
         return (PyObject*)o;
     }
     assert (coro_flags == CO_COROUTINE);
-    PyObject *coro = make_gen(&PyCoro_Type, con, frame);
+    PyObject *coro = make_gen(&PyCoro_Type, func, frame);
     if (!coro) {
         return NULL;
     }
diff --git a/Objects/typeobject.c b/Objects/typeobject.c
index 4180a9d..22e509b 100644
--- a/Objects/typeobject.c
+++ b/Objects/typeobject.c
@@ -8936,7 +8936,7 @@ super_init_without_args(PyFrameObject *f, PyCodeObject *co,
         // "firstarg" is a cell here unless (very unlikely) super()
         // was called from the C-API before the first MAKE_CELL op.
         if (f->f_frame->f_lasti >= 0) {
-            assert(_Py_OPCODE(*co->co_firstinstr) == MAKE_CELL);
+            assert(_Py_OPCODE(*co->co_firstinstr) == MAKE_CELL || _Py_OPCODE(*co->co_firstinstr) == COPY_FREE_VARS);
             assert(PyCell_Check(firstarg));
             firstarg = PyCell_GET(firstarg);
         }