Add function-by-version cache

gvanrossum · Aug 17, 2023 · Aug 8, 2023 · Aug 8, 2023 · Aug 12, 2023 · Aug 12, 2023
commit 138fc0ef7caaffd5885c93c73133e81e60e637fc
diff --git a/Include/internal/pycore_function.h b/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;
+    // Function objects whose func_version % FUNC_VERSION_CACHE_SIZE
+    // once equaled the index in the table. The references are owned:
+    // Call _PyFunction_ClearByVersionCache() to clear.
+    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);
+void _PyFunction_ClearByVersionCache(void);
+
 extern PyObject *_Py_set_function_type_params(
     PyThreadState* unused, PyObject *func, PyObject *type_params);
 

diff --git a/Objects/funcobject.c b/Objects/funcobject.c
@@ -223,7 +223,44 @@ PyFunction_NewWithQualName(PyObject *code, PyObject *globals, PyObject *qualname
     return NULL;
 }
 
-uint32_t _PyFunction_GetVersionForCurrentState(PyFunctionObject *func)
+void
+_PyFunction_SetVersion(PyFunctionObject *func, uint32_t version)
+{
+    func->func_version = version;
+    if (version != 0) {
+        PyInterpreterState *interp = _PyInterpreterState_GET();
+        PyFunctionObject **slot =
+            interp->func_state.func_version_cache
+            + (version % FUNC_VERSION_CACHE_SIZE);
+        Py_XSETREF(*slot, (PyFunctionObject *)Py_NewRef(func));
+    }
+}
+
+PyFunctionObject *
+_PyFunction_LookupByVersion(uint32_t version)
+{
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    PyFunctionObject **slot =
+        interp->func_state.func_version_cache
+        + (version % FUNC_VERSION_CACHE_SIZE);
+    if (*slot != NULL && (*slot)->func_version == version) {
+        return (PyFunctionObject *)Py_NewRef(*slot);
+    }
+    return NULL;
+}
+
+void
+_PyFunction_ClearByVersionCache(void)
+{
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    for (int i = 0; i < FUNC_VERSION_CACHE_SIZE; i++) {
+        PyFunctionObject **slot = interp->func_state.func_version_cache + i;
+        Py_CLEAR(*slot);
+    }
+}
+
+uint32_t
+_PyFunction_GetVersionForCurrentState(PyFunctionObject *func)
 {
     if (func->func_version != 0) {
         return func->func_version;
@@ -236,7 +273,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;
 }
 

@@ -3549,7 +3549,8 @@ dummy_func(
                 goto error;
             }
 
-            func_obj->func_version = ((PyCodeObject *)codeobj)->co_version;
+            _PyFunction_SetVersion(
+                func_obj, ((PyCodeObject *)codeobj)->co_version);
             func = (PyObject *)func_obj;
         }
 

diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
diff --git a/Python/sysmodule.c b/Python/sysmodule.c
@@ -2079,6 +2079,8 @@ sys__clear_type_cache_impl(PyObject *module)
 /*[clinic end generated code: output=20e48ca54a6f6971 input=127f3e04a8d9b555]*/
 {
     PyType_ClearCache();
+    // Also clear the function-by-version cache
+    _PyFunction_ClearByVersionCache();
     Py_RETURN_NONE;
 }