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.

7 files changed, 93 insertions, 62 deletions

diff --git a/Python/bltinmodule.c b/Python/bltinmodule.c
index 9e3b25c..6763f99 100644
--- a/Python/bltinmodule.c
+++ b/Python/bltinmodule.c
@@ -212,9 +212,8 @@ builtin___build_class__(PyObject *self, PyObject *const *args, Py_ssize_t nargs,
                      Py_TYPE(ns)->tp_name);
         goto error;
     }
-    PyFrameConstructor *f =  PyFunction_AS_FRAME_CONSTRUCTOR(func);
     PyThreadState *tstate = _PyThreadState_GET();
-    cell = _PyEval_Vector(tstate, f, ns, NULL, 0, NULL);
+    cell = _PyEval_Vector(tstate, (PyFunctionObject *)func, ns, NULL, 0, NULL);
     if (cell != NULL) {
         if (bases != orig_bases) {
             if (PyMapping_SetItemString(ns, "__orig_bases__", orig_bases) < 0) {
diff --git a/Python/ceval.c b/Python/ceval.c
index 1d69708..0aec5aa 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -14,6 +14,7 @@
 #include "pycore_call.h"          // _PyObject_FastCallDictTstate()
 #include "pycore_ceval.h"         // _PyEval_SignalAsyncExc()
 #include "pycore_code.h"
+#include "pycore_function.h"
 #include "pycore_initconfig.h"    // _PyStatus_OK()
 #include "pycore_long.h"          // _PyLong_GetZero()
 #include "pycore_object.h"        // _PyObject_GC_TRACK()
@@ -98,7 +99,7 @@ static void format_kwargs_error(PyThreadState *, PyObject *func, PyObject *kwarg
 static void format_awaitable_error(PyThreadState *, PyTypeObject *, int, int);
 static int get_exception_handler(PyCodeObject *, int, int*, int*, int*);
 static InterpreterFrame *
-_PyEvalFramePushAndInit(PyThreadState *tstate, PyFrameConstructor *con,
+_PyEvalFramePushAndInit(PyThreadState *tstate, PyFunctionObject *func,
                         PyObject *locals, PyObject* const* args,
                         size_t argcount, PyObject *kwnames);
 static int
@@ -1135,7 +1136,13 @@ PyEval_EvalCode(PyObject *co, PyObject *globals, PyObject *locals)
         .fc_kwdefaults = NULL,
         .fc_closure = NULL
     };
-    return _PyEval_Vector(tstate, &desc, locals, NULL, 0, NULL);
+    PyFunctionObject *func = _PyFunction_FromConstructor(&desc);
+    if (func == NULL) {
+        return NULL;
+    }
+    PyObject *res = _PyEval_Vector(tstate, func, locals, NULL, 0, NULL);
+    Py_DECREF(func);
+    return res;
 }
 
 
@@ -1570,7 +1577,7 @@ trace_function_entry(PyThreadState *tstate, InterpreterFrame *frame)
 }
 
 static PyObject *
-make_coro(PyThreadState *tstate, PyFrameConstructor *con,
+make_coro(PyThreadState *tstate, PyFunctionObject *func,
           PyObject *locals,
           PyObject* const* args, size_t argcount,
           PyObject *kwnames);
@@ -2240,7 +2247,7 @@ check_eval_breaker:
             if (new_frame == NULL) {
                 goto error;
             }
-            _PyFrame_InitializeSpecials(new_frame, PyFunction_AS_FRAME_CONSTRUCTOR(getitem),
+            _PyFrame_InitializeSpecials(new_frame, getitem,
                                         NULL, code->co_nlocalsplus);
             STACK_SHRINK(2);
             new_frame->localsplus[0] = container;
@@ -3179,6 +3186,20 @@ check_eval_breaker:
             DISPATCH();
         }
 
+        TARGET(COPY_FREE_VARS) {
+            /* Copy closure variables to free variables */
+            PyCodeObject *co = frame->f_code;
+            PyObject *closure = frame->f_func->func_closure;
+            int offset = co->co_nlocals + co->co_nplaincellvars;
+            assert(oparg == co->co_nfreevars);
+            for (int i = 0; i < oparg; ++i) {
+                PyObject *o = PyTuple_GET_ITEM(closure, i);
+                Py_INCREF(o);
+                frame->localsplus[offset + i] = o;
+            }
+            DISPATCH();
+        }
+
         TARGET(BUILD_STRING) {
             PyObject *str;
             PyObject *empty = PyUnicode_New(0, 0);
@@ -4423,9 +4444,9 @@ check_eval_breaker:
                     PyObject *locals = code_flags & CO_OPTIMIZED ? NULL : PyFunction_GET_GLOBALS(function);
                     STACK_SHRINK(oparg);
                     InterpreterFrame *new_frame = _PyEvalFramePushAndInit(
-                        tstate, PyFunction_AS_FRAME_CONSTRUCTOR(function), locals,
-                                                                stack_pointer,
-                                                                nargs, kwnames);
+                        tstate, (PyFunctionObject *)function, locals,
+                        stack_pointer, nargs, kwnames
+                    );
                     STACK_SHRINK(postcall_shrink);
                     // The frame has stolen all the arguments from the stack,
                     // so there is no need to clean them up.
@@ -4506,7 +4527,7 @@ check_eval_breaker:
             if (new_frame == NULL) {
                 goto error;
             }
-            _PyFrame_InitializeSpecials(new_frame, PyFunction_AS_FRAME_CONSTRUCTOR(func),
+            _PyFrame_InitializeSpecials(new_frame, func,
                                         NULL, code->co_nlocalsplus);
             STACK_SHRINK(argcount);
             for (int i = 0; i < argcount; i++) {
@@ -5426,11 +5447,11 @@ get_exception_handler(PyCodeObject *code, int index, int *level, int *handler, i
 }
 
 static int
-initialize_locals(PyThreadState *tstate, PyFrameConstructor *con,
+initialize_locals(PyThreadState *tstate, PyFunctionObject *func,
     PyObject **localsplus, PyObject *const *args,
     Py_ssize_t argcount, PyObject *kwnames)
 {
-    PyCodeObject *co = (PyCodeObject*)con->fc_code;
+    PyCodeObject *co = (PyCodeObject*)func->func_code;
     const Py_ssize_t total_args = co->co_argcount + co->co_kwonlyargcount;
 
     /* Create a dictionary for keyword parameters (**kwags) */
@@ -5495,7 +5516,7 @@ initialize_locals(PyThreadState *tstate, PyFrameConstructor *con,
             if (keyword == NULL || !PyUnicode_Check(keyword)) {
                 _PyErr_Format(tstate, PyExc_TypeError,
                             "%U() keywords must be strings",
-                          con->fc_qualname);
+                          func->func_qualname);
                 goto kw_fail;
             }
 
@@ -5527,14 +5548,14 @@ initialize_locals(PyThreadState *tstate, PyFrameConstructor *con,
                 if (co->co_posonlyargcount
                     && positional_only_passed_as_keyword(tstate, co,
                                                         kwcount, kwnames,
-                                                     con->fc_qualname))
+                                                        func->func_qualname))
                 {
                     goto kw_fail;
                 }
 
                 _PyErr_Format(tstate, PyExc_TypeError,
                             "%U() got an unexpected keyword argument '%S'",
-                          con->fc_qualname, keyword);
+                          func->func_qualname, keyword);
                 goto kw_fail;
             }
 
@@ -5555,7 +5576,7 @@ initialize_locals(PyThreadState *tstate, PyFrameConstructor *con,
             if (localsplus[j] != NULL) {
                 _PyErr_Format(tstate, PyExc_TypeError,
                             "%U() got multiple values for argument '%S'",
-                          con->fc_qualname, keyword);
+                          func->func_qualname, keyword);
                 goto kw_fail;
             }
             localsplus[j] = value;
@@ -5564,14 +5585,14 @@ initialize_locals(PyThreadState *tstate, PyFrameConstructor *con,
 
     /* Check the number of positional arguments */
     if ((argcount > co->co_argcount) && !(co->co_flags & CO_VARARGS)) {
-        too_many_positional(tstate, co, argcount, con->fc_defaults, localsplus,
-                            con->fc_qualname);
+        too_many_positional(tstate, co, argcount, func->func_defaults, localsplus,
+                            func->func_qualname);
         goto fail_post_args;
     }
 
     /* Add missing positional arguments (copy default values from defs) */
     if (argcount < co->co_argcount) {
-        Py_ssize_t defcount = con->fc_defaults == NULL ? 0 : PyTuple_GET_SIZE(con->fc_defaults);
+        Py_ssize_t defcount = func->func_defaults == NULL ? 0 : PyTuple_GET_SIZE(func->func_defaults);
         Py_ssize_t m = co->co_argcount - defcount;
         Py_ssize_t missing = 0;
         for (i = argcount; i < m; i++) {
@@ -5581,7 +5602,7 @@ initialize_locals(PyThreadState *tstate, PyFrameConstructor *con,
         }
         if (missing) {
             missing_arguments(tstate, co, missing, defcount, localsplus,
-                              con->fc_qualname);
+                              func->func_qualname);
             goto fail_post_args;
         }
         if (n > m)
@@ -5589,7 +5610,7 @@ initialize_locals(PyThreadState *tstate, PyFrameConstructor *con,
         else
             i = 0;
         if (defcount) {
-            PyObject **defs = &PyTuple_GET_ITEM(con->fc_defaults, 0);
+            PyObject **defs = &PyTuple_GET_ITEM(func->func_defaults, 0);
             for (; i < defcount; i++) {
                 if (localsplus[m+i] == NULL) {
                     PyObject *def = defs[i];
@@ -5607,8 +5628,8 @@ initialize_locals(PyThreadState *tstate, PyFrameConstructor *con,
             if (localsplus[i] != NULL)
                 continue;
             PyObject *varname = PyTuple_GET_ITEM(co->co_localsplusnames, i);
-            if (con->fc_kwdefaults != NULL) {
-                PyObject *def = PyDict_GetItemWithError(con->fc_kwdefaults, varname);
+            if (func->func_kwdefaults != NULL) {
+                PyObject *def = PyDict_GetItemWithError(func->func_kwdefaults, varname);
                 if (def) {
                     Py_INCREF(def);
                     localsplus[i] = def;
@@ -5622,16 +5643,10 @@ initialize_locals(PyThreadState *tstate, PyFrameConstructor *con,
         }
         if (missing) {
             missing_arguments(tstate, co, missing, -1, localsplus,
-                              con->fc_qualname);
+                              func->func_qualname);
             goto fail_post_args;
         }
     }
-    /* Copy closure variables to free variables */
-    for (i = 0; i < co->co_nfreevars; ++i) {
-        PyObject *o = PyTuple_GET_ITEM(con->fc_closure, i);
-        Py_INCREF(o);
-        localsplus[co->co_nlocals + co->co_nplaincellvars + i] = o;
-    }
     return 0;
 
 fail_pre_positional:
@@ -5653,24 +5668,24 @@ fail_post_args:
 
 static InterpreterFrame *
 make_coro_frame(PyThreadState *tstate,
-           PyFrameConstructor *con, PyObject *locals,
+           PyFunctionObject *func, PyObject *locals,
            PyObject *const *args, Py_ssize_t argcount,
            PyObject *kwnames)
 {
     assert(is_tstate_valid(tstate));
-    assert(con->fc_defaults == NULL || PyTuple_CheckExact(con->fc_defaults));
-    PyCodeObject *code = (PyCodeObject *)con->fc_code;
+    assert(func->func_defaults == NULL || PyTuple_CheckExact(func->func_defaults));
+    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) {
         goto fail_no_memory;
     }
-    _PyFrame_InitializeSpecials(frame, con, locals, code->co_nlocalsplus);
+    _PyFrame_InitializeSpecials(frame, func, locals, code->co_nlocalsplus);
     for (int i = 0; i < code->co_nlocalsplus; i++) {
         frame->localsplus[i] = NULL;
     }
     assert(frame->frame_obj == NULL);
-    if (initialize_locals(tstate, con, frame->localsplus, args, argcount, kwnames)) {
+    if (initialize_locals(tstate, func, frame->localsplus, args, argcount, kwnames)) {
         _PyFrame_Clear(frame, 1);
         return NULL;
     }
@@ -5692,17 +5707,17 @@ fail_no_memory:
 
 /* Consumes all the references to the args */
 static PyObject *
-make_coro(PyThreadState *tstate, PyFrameConstructor *con,
+make_coro(PyThreadState *tstate, PyFunctionObject *func,
           PyObject *locals,
           PyObject* const* args, size_t argcount,
           PyObject *kwnames)
 {
-    assert (((PyCodeObject *)con->fc_code)->co_flags & (CO_GENERATOR | CO_COROUTINE | CO_ASYNC_GENERATOR));
-    InterpreterFrame *frame = make_coro_frame(tstate, con, locals, args, argcount, kwnames);
+    assert (((PyCodeObject *)func->func_code)->co_flags & (CO_GENERATOR | CO_COROUTINE | CO_ASYNC_GENERATOR));
+    InterpreterFrame *frame = make_coro_frame(tstate, func, locals, args, argcount, kwnames);
     if (frame == NULL) {
         return NULL;
     }
-    PyObject *gen = _Py_MakeCoro(con, frame);
+    PyObject *gen = _Py_MakeCoro(func, frame);
     if (gen == NULL) {
         return NULL;
     }
@@ -5711,22 +5726,22 @@ make_coro(PyThreadState *tstate, PyFrameConstructor *con,
 
 /* Consumes all the references to the args */
 static InterpreterFrame *
-_PyEvalFramePushAndInit(PyThreadState *tstate, PyFrameConstructor *con,
+_PyEvalFramePushAndInit(PyThreadState *tstate, PyFunctionObject *func,
                         PyObject *locals, PyObject* const* args,
                         size_t argcount, PyObject *kwnames)
 {
-    PyCodeObject * code = (PyCodeObject *)con->fc_code;
+    PyCodeObject * code = (PyCodeObject *)func->func_code;
     size_t size = code->co_nlocalsplus + code->co_stacksize + FRAME_SPECIALS_SIZE;
     InterpreterFrame *frame = _PyThreadState_BumpFramePointer(tstate, size);
     if (frame == NULL) {
         goto fail;
     }
-    _PyFrame_InitializeSpecials(frame, con, locals, code->co_nlocalsplus);
+    _PyFrame_InitializeSpecials(frame, func, locals, code->co_nlocalsplus);
     PyObject **localsarray = &frame->localsplus[0];
     for (int i = 0; i < code->co_nlocalsplus; i++) {
         localsarray[i] = NULL;
     }
-    if (initialize_locals(tstate, con, localsarray, args, argcount, kwnames)) {
+    if (initialize_locals(tstate, func, localsarray, args, argcount, kwnames)) {
         _PyFrame_Clear(frame, 0);
         return NULL;
     }
@@ -5761,12 +5776,12 @@ _PyEvalFrameClearAndPop(PyThreadState *tstate, InterpreterFrame * frame)
 }
 
 PyObject *
-_PyEval_Vector(PyThreadState *tstate, PyFrameConstructor *con,
+_PyEval_Vector(PyThreadState *tstate, PyFunctionObject *func,
                PyObject *locals,
                PyObject* const* args, size_t argcount,
                PyObject *kwnames)
 {
-    PyCodeObject *code = (PyCodeObject *)con->fc_code;
+    PyCodeObject *code = (PyCodeObject *)func->func_code;
     /* _PyEvalFramePushAndInit and make_coro consume
      * all the references to their arguments
      */
@@ -5782,10 +5797,10 @@ _PyEval_Vector(PyThreadState *tstate, PyFrameConstructor *con,
     int is_coro = code->co_flags &
         (CO_GENERATOR | CO_COROUTINE | CO_ASYNC_GENERATOR);
     if (is_coro) {
-        return make_coro(tstate, con, locals, args, argcount, kwnames);
+        return make_coro(tstate, func, locals, args, argcount, kwnames);
     }
     InterpreterFrame *frame = _PyEvalFramePushAndInit(
-        tstate, con, locals, args, argcount, kwnames);
+        tstate, func, locals, args, argcount, kwnames);
     if (frame == NULL) {
         return NULL;
     }
@@ -5869,9 +5884,14 @@ PyEval_EvalCodeEx(PyObject *_co, PyObject *globals, PyObject *locals,
         .fc_kwdefaults = kwdefs,
         .fc_closure = closure
     };
-    res = _PyEval_Vector(tstate, &constr, locals,
+    PyFunctionObject *func = _PyFunction_FromConstructor(&constr);
+    if (func == NULL) {
+        return NULL;
+    }
+    res = _PyEval_Vector(tstate, func, locals,
                          allargs, argcount,
                          kwnames);
+    Py_DECREF(func);
     if (kwcount) {
         Py_DECREF(kwnames);
         PyMem_Free(newargs);
diff --git a/Python/compile.c b/Python/compile.c
index 40bd1fd..87de7ba 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -1171,6 +1171,7 @@ stack_effect(int opcode, int oparg, int jump)
 
         /* Closures */
         case MAKE_CELL:
+        case COPY_FREE_VARS:
             return 0;
         case LOAD_CLOSURE:
             return 1;
@@ -7611,7 +7612,7 @@ insert_instruction(basicblock *block, int pos, struct instr *instr) {
 
 static int
 insert_prefix_instructions(struct compiler *c, basicblock *entryblock,
-                           int *fixed)
+                           int *fixed, int nfreevars)
 {
 
     int flags = compute_code_flags(c);
@@ -7684,6 +7685,22 @@ insert_prefix_instructions(struct compiler *c, basicblock *entryblock,
         }
     }
 
+    if (nfreevars) {
+        struct instr copy_frees = {
+            .i_opcode = COPY_FREE_VARS,
+            .i_oparg = nfreevars,
+            .i_lineno = -1,
+            .i_col_offset = -1,
+            .i_end_lineno = -1,
+            .i_end_col_offset = -1,
+            .i_target = NULL,
+        };
+        if (insert_instruction(entryblock, 0, &copy_frees) < 0) {
+            return -1;
+        }
+
+    }
+
     return 0;
 }
 
@@ -7818,7 +7835,7 @@ assemble(struct compiler *c, int addNone)
     }
 
     // This must be called before fix_cell_offsets().
-    if (insert_prefix_instructions(c, entryblock, cellfixedoffsets)) {
+    if (insert_prefix_instructions(c, entryblock, cellfixedoffsets, nfreevars)) {
         goto error;
     }
 
diff --git a/Python/frame.c b/Python/frame.c
index a5c93ea..79b0f77 100644
--- a/Python/frame.c
+++ b/Python/frame.c
@@ -8,9 +8,8 @@ int
 _PyFrame_Traverse(InterpreterFrame *frame, visitproc visit, void *arg)
 {
     Py_VISIT(frame->frame_obj);
-    Py_VISIT(frame->f_globals);
-    Py_VISIT(frame->f_builtins);
     Py_VISIT(frame->f_locals);
+    Py_VISIT(frame->f_func);
     Py_VISIT(frame->f_code);
    /* locals */
     PyObject **locals = _PyFrame_GetLocalsArray(frame);
@@ -62,8 +61,7 @@ clear_specials(InterpreterFrame *frame)
     frame->generator = NULL;
     Py_XDECREF(frame->frame_obj);
     Py_XDECREF(frame->f_locals);
-    Py_DECREF(frame->f_globals);
-    Py_DECREF(frame->f_builtins);
+    Py_DECREF(frame->f_func);
     Py_DECREF(frame->f_code);
 }
 
diff --git a/Python/opcode_targets.h b/Python/opcode_targets.h
index 903b967..c9d430d 100644
--- a/Python/opcode_targets.h
+++ b/Python/opcode_targets.h
@@ -148,7 +148,7 @@ static void *opcode_targets[256] = {
     &&TARGET_SET_ADD,
     &&TARGET_MAP_ADD,
     &&TARGET_LOAD_CLASSDEREF,
-    &&_unknown_opcode,
+    &&TARGET_COPY_FREE_VARS,
     &&_unknown_opcode,
     &&_unknown_opcode,
     &&TARGET_MATCH_CLASS,
diff --git a/Python/pystate.c b/Python/pystate.c
index 273982b..56db095d 100644
--- a/Python/pystate.c
+++ b/Python/pystate.c
@@ -2087,9 +2087,9 @@ _PyThreadState_BumpFramePointerSlow(PyThreadState *tstate, size_t size)
 
 
 InterpreterFrame *
-_PyThreadState_PushFrame(PyThreadState *tstate, PyFrameConstructor *con, PyObject *locals)
+_PyThreadState_PushFrame(PyThreadState *tstate, PyFunctionObject *func, PyObject *locals)
 {
-    PyCodeObject *code = (PyCodeObject *)con->fc_code;
+    PyCodeObject *code = (PyCodeObject *)func->func_code;
     int nlocalsplus = code->co_nlocalsplus;
     size_t size = nlocalsplus + code->co_stacksize +
         FRAME_SPECIALS_SIZE;
@@ -2097,7 +2097,7 @@ _PyThreadState_PushFrame(PyThreadState *tstate, PyFrameConstructor *con, PyObjec
     if (frame == NULL) {
         return NULL;
     }
-    _PyFrame_InitializeSpecials(frame, con, locals, nlocalsplus);
+    _PyFrame_InitializeSpecials(frame, func, locals, nlocalsplus);
     for (int i=0; i < nlocalsplus; i++) {
         frame->localsplus[i] = NULL;
     }
diff --git a/Python/specialize.c b/Python/specialize.c
index 130da00..f5f1213 100644
--- a/Python/specialize.c
+++ b/Python/specialize.c
@@ -479,7 +479,7 @@ initial_counter_value(void) {
 #define SPEC_FAIL_WRONG_NUMBER_ARGUMENTS 9
 #define SPEC_FAIL_CO_NOT_OPTIMIZED 10
 /* SPEC_FAIL_METHOD  defined as 11 above */
-#define SPEC_FAIL_FREE_VARS 12
+
 #define SPEC_FAIL_PYCFUNCTION 13
 #define SPEC_FAIL_PYCFUNCTION_WITH_KEYWORDS 14
 #define SPEC_FAIL_PYCFUNCTION_FAST_WITH_KEYWORDS 15
@@ -1158,9 +1158,6 @@ function_kind(PyCodeObject *code) {
     if ((flags & CO_OPTIMIZED) == 0) {
         return SPEC_FAIL_CO_NOT_OPTIMIZED;
     }
-    if (code->co_nfreevars) {
-        return SPEC_FAIL_FREE_VARS;
-    }
     return SIMPLE_FUNCTION;
 }