Support RETURN from tracing through a function

Include/internal/pycore_ceval.h

@@ -171,6 +171,7 @@ void _PyEval_FormatKwargsError(PyThreadState *tstate, PyObject *func, PyObject *
 PyObject *_PyEval_MatchClass(PyThreadState *tstate, PyObject *subject, PyObject *type, Py_ssize_t nargs, PyObject *kwargs);
 PyObject *_PyEval_MatchKeys(PyThreadState *tstate, PyObject *map, PyObject *keys);
 int _PyEval_UnpackIterable(PyThreadState *tstate, PyObject *v, int argcnt, int argcntafter, PyObject **sp);
+void _PyEval_FrameClearAndPop(PyThreadState *tstate, _PyInterpreterFrame *frame);
 
 
 #ifdef __cplusplus

Include/internal/pycore_function.h

@@ -16,13 +16,22 @@ extern PyObject* _PyFunction_Vectorcall(
 
 #define FUNC_MAX_WATCHERS 8
 
+#define FUNC_VERSION_CACHE_SIZE (1<<12)  /* Must be a power of 2 */
 struct _py_func_state {
     uint32_t next_version;
+    // Borrowed references to function objects whose
+    // func_version % FUNC_VERSION_CACHE_SIZE
+    // once was equal to the index in the table.
+    // They are cleared when the function is deallocated.
+    PyFunctionObject *func_version_cache[FUNC_VERSION_CACHE_SIZE];
 };
 
 extern PyFunctionObject* _PyFunction_FromConstructor(PyFrameConstructor *constr);
 
 extern uint32_t _PyFunction_GetVersionForCurrentState(PyFunctionObject *func);
+extern void _PyFunction_SetVersion(PyFunctionObject *func, uint32_t version);
+PyFunctionObject *_PyFunction_LookupByVersion(uint32_t version);
+
 extern PyObject *_Py_set_function_type_params(
     PyThreadState* unused, PyObject *func, PyObject *type_params);
 

Lib/test/test_capi/test_misc.py

@@ -2618,6 +2618,24 @@ def testfunc(it):
             with self.assertRaises(StopIteration):
                 next(it)
 
+    def test_call_py_exact_args(self):
+        def testfunc(n):
+            def dummy(x):
+                return x+1
+            for i in range(n):
+                dummy(i)
+
+        opt = _testinternalcapi.get_uop_optimizer()
+        with temporary_optimizer(opt):
+            testfunc(10)
+
+        ex = get_first_executor(testfunc)
+        self.assertIsNotNone(ex)
+        uops = {opname for opname, _, _ in ex}
+        self.assertIn("_PUSH_FRAME", uops)
+        self.assertIn("_BINARY_OP_ADD_INT", uops)
+
+
 
 if __name__ == "__main__":
     unittest.main()

Lib/test/test_code.py

@@ -264,7 +264,7 @@ def func2():
             ("co_posonlyargcount", 0),
             ("co_kwonlyargcount", 0),
             ("co_nlocals", 1),
-            ("co_stacksize", 0),
+            ("co_stacksize", 1),
             ("co_flags", code.co_flags | inspect.CO_COROUTINE),
             ("co_firstlineno", 100),
             ("co_code", code2.co_code),

Objects/codeobject.c

@@ -395,6 +395,9 @@ init_code(PyCodeObject *co, struct _PyCodeConstructor *con)
     int nlocals, ncellvars, nfreevars;
     get_localsplus_counts(con->localsplusnames, con->localspluskinds,
                           &nlocals, &ncellvars, &nfreevars);
+    if (con->stacksize == 0) {
+        con->stacksize = 1;
+    }
 
     co->co_filename = Py_NewRef(con->filename);
     co->co_name = Py_NewRef(con->name);

Objects/funcobject.c

@@ -223,7 +223,73 @@ PyFunction_NewWithQualName(PyObject *code, PyObject *globals, PyObject *qualname
     return NULL;
 }
 
-uint32_t _PyFunction_GetVersionForCurrentState(PyFunctionObject *func)
+/*
+Function versions
+-----------------
+
+Function versions are used to detect when a function object has been
+updated, invalidating inline cache data used by the `CALL` bytecode
+(notably `CALL_PY_EXACT_ARGS` and a few other `CALL` specializations).
+
+They are also used by the Tier 2 superblock creation code to find
+the function being called (and from there the code object).
+
+How does a function's `func_version` field get initialized?
+
+- `PyFunction_New` and friends initialize it to 0.
+- The `MAKE_FUNCTION` instruction sets it from the code's `co_version`.
+- It is reset to 0 when various attributes like `__code__` are set.
+- A new version is allocated by `_PyFunction_GetVersionForCurrentState`
+  when the specializer needs a version and the version is 0.
+
+The latter allocates versions using a counter in the interpreter state;
+when the counter wraps around to 0, no more versions are allocated.
+There is one other special case: functions with a non-standard
+`vectorcall` field are not given a version.
+
+When the function version is 0, the `CALL` bytecode is not specialized.
+
+Code object versions
+--------------------
+
+So where to code objects get their `co_version`? There is a single
+static global counter, `_Py_next_func_version`. This is initialized in
+the generated (!) file `Python/deepfreeze/deepfreeze.c`, to 1 plus the
+number of deep-frozen function objects in that file.
+(In `_bootstrap_python.c` and `freeze_module.c` it is initialized to 1.)
+
+Code objects get a new `co_version` allocated from this counter upon
+creation. Since code objects are nominally immutable, `co_version` can
+not be invalidated. The only way it can be 0 is when 2**32 or more
+code objects have been created during the process's lifetime.
+(The counter isn't reset by `fork()`, extending the lifetime.)
+*/
+
+void
+_PyFunction_SetVersion(PyFunctionObject *func, uint32_t version)
+{
+    func->func_version = version;
+    if (version != 0) {
+        PyInterpreterState *interp = _PyInterpreterState_GET();
+        interp->func_state.func_version_cache[
+            version % FUNC_VERSION_CACHE_SIZE] = func;
+    }
+}
+
+PyFunctionObject *
+_PyFunction_LookupByVersion(uint32_t version)
+{
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    PyFunctionObject *func = interp->func_state.func_version_cache[
+        version % FUNC_VERSION_CACHE_SIZE];
+    if (func != NULL && func->func_version == version) {
+        return (PyFunctionObject *)Py_NewRef(func);
+    }
+    return NULL;
+}
+
+uint32_t
+_PyFunction_GetVersionForCurrentState(PyFunctionObject *func)
 {
     if (func->func_version != 0) {
         return func->func_version;
@@ -236,7 +302,7 @@ uint32_t _PyFunction_GetVersionForCurrentState(PyFunctionObject *func)
         return 0;
     }
     uint32_t v = interp->func_state.next_version++;
-    func->func_version = v;
+    _PyFunction_SetVersion(func, v);
     return v;
 }
 
@@ -851,6 +917,15 @@ func_dealloc(PyFunctionObject *op)
     if (op->func_weakreflist != NULL) {
         PyObject_ClearWeakRefs((PyObject *) op);
     }
+    if (op->func_version != 0) {
+        PyInterpreterState *interp = _PyInterpreterState_GET();
+        PyFunctionObject **slot =
+            interp->func_state.func_version_cache
+            + (op->func_version % FUNC_VERSION_CACHE_SIZE);
+        if (*slot == op) {
+            *slot = NULL;
+        }
+    }
     (void)func_clear(op);
     // These aren't cleared by func_clear().
     Py_DECREF(op->func_code);