GH-91049: Introduce set vectorcall field API for PyFunctionObject (GH-92257)

Co-authored-by: Andrew Frost <adfrost@fb.com>
Co-authored-by: Itamar Ostricher <itamarost@gmail.com>

9 files changed, 115 insertions, 3 deletions

diff --git a/Doc/c-api/function.rst b/Doc/c-api/function.rst
index 56c1839..df88e85 100644
--- a/Doc/c-api/function.rst
+++ b/Doc/c-api/function.rst
@@ -83,6 +83,15 @@ There are a few functions specific to Python functions.
    Raises :exc:`SystemError` and returns ``-1`` on failure.
 
 
+.. c:function:: void PyFunction_SetVectorcall(PyFunctionObject *func, vectorcallfunc vectorcall)
+
+   Set the vectorcall field of a given function object *func*.
+
+   Warning: extensions using this API must preserve the behavior
+   of the unaltered (default) vectorcall function!
+
+   .. versionadded:: 3.12
+
 .. c:function:: PyObject* PyFunction_GetClosure(PyObject *op)
 
    Return the closure associated with the function object *op*. This can be ``NULL``
diff --git a/Doc/whatsnew/3.12.rst b/Doc/whatsnew/3.12.rst
index 4a6b9c0..90355a7 100644
--- a/Doc/whatsnew/3.12.rst
+++ b/Doc/whatsnew/3.12.rst
@@ -493,6 +493,10 @@ New Features
   functions in all running threads in addition to the calling one. (Contributed
   by Pablo Galindo in :gh:`93503`.)
 
+* Added new function :c:func:`PyFunction_SetVectorcall` to the C API
+  which sets the vectorcall field of a given :c:type:`PyFunctionObject`.
+  (Contributed by Andrew Frost in :gh:`92257`.)
+
 Porting to Python 3.12
 ----------------------
 
diff --git a/Include/cpython/funcobject.h b/Include/cpython/funcobject.h
index 05f7665..dd8f20b 100644
--- a/Include/cpython/funcobject.h
+++ b/Include/cpython/funcobject.h
@@ -48,7 +48,8 @@ typedef struct {
      *     defaults
      *     kwdefaults (only if the object changes, not the contents of the dict)
      *     code
-     *     annotations */
+     *     annotations
+     *     vectorcall function pointer */
     uint32_t func_version;
 
     /* Invariant:
@@ -69,6 +70,7 @@ PyAPI_FUNC(PyObject *) PyFunction_GetGlobals(PyObject *);
 PyAPI_FUNC(PyObject *) PyFunction_GetModule(PyObject *);
 PyAPI_FUNC(PyObject *) PyFunction_GetDefaults(PyObject *);
 PyAPI_FUNC(int) PyFunction_SetDefaults(PyObject *, PyObject *);
+PyAPI_FUNC(void) PyFunction_SetVectorcall(PyFunctionObject *, vectorcallfunc);
 PyAPI_FUNC(PyObject *) PyFunction_GetKwDefaults(PyObject *);
 PyAPI_FUNC(int) PyFunction_SetKwDefaults(PyObject *, PyObject *);
 PyAPI_FUNC(PyObject *) PyFunction_GetClosure(PyObject *);
diff --git a/Lib/test/test_call.py b/Lib/test/test_call.py
index c00de27..c1a3862 100644
--- a/Lib/test/test_call.py
+++ b/Lib/test/test_call.py
@@ -580,6 +580,9 @@ def testfunction_kw(self, *, kw):
     return self
 
 
+QUICKENING_WARMUP_DELAY = 8
+
+
 class TestPEP590(unittest.TestCase):
 
     def test_method_descriptor_flag(self):
@@ -760,6 +763,54 @@ class TestPEP590(unittest.TestCase):
                 self.assertEqual(expected, meth(*args1, **kwargs))
                 self.assertEqual(expected, wrapped(*args, **kwargs))
 
+    def test_setvectorcall(self):
+        from _testcapi import function_setvectorcall
+        def f(num): return num + 1
+        assert_equal = self.assertEqual
+        num = 10
+        assert_equal(11, f(num))
+        function_setvectorcall(f)
+        # make sure specializer is triggered by running > 50 times
+        for _ in range(10 * QUICKENING_WARMUP_DELAY):
+            assert_equal("overridden", f(num))
+
+    def test_setvectorcall_load_attr_specialization_skip(self):
+        from _testcapi import function_setvectorcall
+
+        class X:
+            def __getattribute__(self, attr):
+                return attr
+
+        assert_equal = self.assertEqual
+        x = X()
+        assert_equal("a", x.a)
+        function_setvectorcall(X.__getattribute__)
+        # make sure specialization doesn't trigger
+        # when vectorcall is overridden
+        for _ in range(QUICKENING_WARMUP_DELAY):
+            assert_equal("overridden", x.a)
+
+    def test_setvectorcall_load_attr_specialization_deopt(self):
+        from _testcapi import function_setvectorcall
+
+        class X:
+            def __getattribute__(self, attr):
+                return attr
+
+        def get_a(x):
+            return x.a
+
+        assert_equal = self.assertEqual
+        x = X()
+        # trigger LOAD_ATTR_GETATTRIBUTE_OVERRIDDEN specialization
+        for _ in range(QUICKENING_WARMUP_DELAY):
+            assert_equal("a", get_a(x))
+        function_setvectorcall(X.__getattribute__)
+        # make sure specialized LOAD_ATTR_GETATTRIBUTE_OVERRIDDEN
+        # gets deopted due to overridden vectorcall
+        for _ in range(QUICKENING_WARMUP_DELAY):
+            assert_equal("overridden", get_a(x))
+
     @requires_limited_api
     def test_vectorcall_limited(self):
         from _testcapi import pyobject_vectorcall
diff --git a/Misc/NEWS.d/next/C API/2022-05-03-19-35-37.gh-issue-92193.61VoFL.rst b/Misc/NEWS.d/next/C API/2022-05-03-19-35-37.gh-issue-92193.61VoFL.rst
new file mode 100644
index 0000000..e0755bb
--- /dev/null
+++ b/Misc/NEWS.d/next/C API/2022-05-03-19-35-37.gh-issue-92193.61VoFL.rst
@@ -0,0 +1,5 @@
+Add new function :c:func:`PyFunction_SetVectorcall` to the C API
+which sets the vectorcall field of a given :c:type:`PyFunctionObject`.
+
+Warning: extensions using this API must preserve the behavior
+of the unaltered function!
diff --git a/Modules/_testcapi/vectorcall.c b/Modules/_testcapi/vectorcall.c
index 626706e..e9c863a 100644
--- a/Modules/_testcapi/vectorcall.c
+++ b/Modules/_testcapi/vectorcall.c
@@ -103,6 +103,24 @@ test_pyobject_vectorcall(PyObject *self, PyObject *args)
 }
 
 static PyObject *
+override_vectorcall(PyObject *callable, PyObject *const *args, size_t nargsf,
+                    PyObject *kwnames)
+{
+    return PyUnicode_FromString("overridden");
+}
+
+static PyObject *
+function_setvectorcall(PyObject *self, PyObject *func)
+{
+    if (!PyFunction_Check(func)) {
+        PyErr_SetString(PyExc_TypeError, "'func' must be a function");
+        return NULL;
+    }
+    PyFunction_SetVectorcall((PyFunctionObject *)func, (vectorcallfunc)override_vectorcall);
+    Py_RETURN_NONE;
+}
+
+static PyObject *
 test_pyvectorcall_call(PyObject *self, PyObject *args)
 {
     PyObject *func;
@@ -244,6 +262,7 @@ static PyMethodDef TestMethods[] = {
     {"pyobject_fastcall", test_pyobject_fastcall, METH_VARARGS},
     {"pyobject_fastcalldict", test_pyobject_fastcalldict, METH_VARARGS},
     {"pyobject_vectorcall", test_pyobject_vectorcall, METH_VARARGS},
+    {"function_setvectorcall", function_setvectorcall, METH_O},
     {"pyvectorcall_call", test_pyvectorcall_call, METH_VARARGS},
     _TESTCAPI_MAKE_VECTORCALL_CLASS_METHODDEF
     _TESTCAPI_HAS_VECTORCALL_FLAG_METHODDEF
diff --git a/Objects/funcobject.c b/Objects/funcobject.c
index 32b4155..7f257a9 100644
--- a/Objects/funcobject.c
+++ b/Objects/funcobject.c
@@ -134,6 +134,9 @@ uint32_t _PyFunction_GetVersionForCurrentState(PyFunctionObject *func)
     if (func->func_version != 0) {
         return func->func_version;
     }
+    if (func->vectorcall != _PyFunction_Vectorcall) {
+        return 0;
+    }
     if (next_func_version == 0) {
         return 0;
     }
@@ -209,6 +212,14 @@ PyFunction_SetDefaults(PyObject *op, PyObject *defaults)
     return 0;
 }
 
+void
+PyFunction_SetVectorcall(PyFunctionObject *func, vectorcallfunc vectorcall)
+{
+    assert(func != NULL);
+    func->func_version = 0;
+    func->vectorcall = vectorcall;
+}
+
 PyObject *
 PyFunction_GetKwDefaults(PyObject *op)
 {
diff --git a/Python/ceval.c b/Python/ceval.c
index b61cc08..8891d6c 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -3126,8 +3126,11 @@ handle_eval_breaker:
             PyObject *getattribute = read_obj(cache->descr);
             assert(Py_IS_TYPE(getattribute, &PyFunction_Type));
             PyFunctionObject *f = (PyFunctionObject *)getattribute;
+            uint32_t func_version = read_u32(cache->keys_version);
+            assert(func_version != 0);
+            DEOPT_IF(f->func_version != func_version, LOAD_ATTR);
             PyCodeObject *code = (PyCodeObject *)f->func_code;
-            DEOPT_IF(((PyCodeObject *)f->func_code)->co_argcount != 2, LOAD_ATTR);
+            assert(code->co_argcount == 2);
             DEOPT_IF(!_PyThreadState_HasStackSpace(tstate, code->co_framesize), CALL);
             STAT_INC(LOAD_ATTR, hit);
 
@@ -4133,7 +4136,10 @@ handle_eval_breaker:
             function = PEEK(total_args + 1);
             int positional_args = total_args - KWNAMES_LEN();
             // Check if the call can be inlined or not
-            if (Py_TYPE(function) == &PyFunction_Type && tstate->interp->eval_frame == NULL) {
+            if (Py_TYPE(function) == &PyFunction_Type &&
+                tstate->interp->eval_frame == NULL &&
+                ((PyFunctionObject *)function)->vectorcall == _PyFunction_Vectorcall)
+            {
                 int code_flags = ((PyCodeObject*)PyFunction_GET_CODE(function))->co_flags;
                 PyObject *locals = code_flags & CO_OPTIMIZED ? NULL : Py_NewRef(PyFunction_GET_GLOBALS(function));
                 STACK_SHRINK(total_args);
diff --git a/Python/specialize.c b/Python/specialize.c
index 93f1d28..b7c321e 100644
--- a/Python/specialize.c
+++ b/Python/specialize.c
@@ -837,6 +837,11 @@ _Py_Specialize_LoadAttr(PyObject *owner, _Py_CODEUNIT *instr, PyObject *name)
             if (!function_check_args(descr, 2, LOAD_ATTR)) {
                 goto fail;
             }
+            uint32_t version = function_get_version(descr, LOAD_ATTR);
+            if (version == 0) {
+                goto fail;
+            }
+            write_u32(lm_cache->keys_version, version);
             /* borrowed */
             write_obj(lm_cache->descr, descr);
             write_u32(lm_cache->type_version, type->tp_version_tag);