gh-89013: Improve the performance of methodcaller (lazy version) (gh-107201)

Misc/NEWS.d/next/Library/2021-08-16-17-52-26.bpo-44850.r8jx5u.rst

@@ -0,0 +1,2 @@
+Improve performance of :func:`operator.methodcaller` using the :pep:`590` ``vectorcall`` convention.
+Patch by Anthony Lee and Pieter Eendebak.

Modules/_operator.c

@@ -1549,10 +1549,77 @@ static PyType_Spec attrgetter_type_spec = {
 typedef struct {
     PyObject_HEAD
     PyObject *name;
-    PyObject *args;
+    PyObject *xargs; // reference to arguments passed in constructor
     PyObject *kwds;
+    PyObject **vectorcall_args;  /* Borrowed references */
+    PyObject *vectorcall_kwnames;
+    vectorcallfunc vectorcall;
 } methodcallerobject;
 
+static int _methodcaller_initialize_vectorcall(methodcallerobject* mc)
+{
+    PyObject* args = mc->xargs;
+    PyObject* kwds = mc->kwds;
+
+    Py_ssize_t nargs = PyTuple_GET_SIZE(args);
+    assert(nargs > 0);
+    mc->vectorcall_args = PyMem_Calloc(
+        nargs + (kwds ? PyDict_Size(kwds) : 0),
+        sizeof(PyObject*));
+    if (!mc->vectorcall_args) {
+        PyErr_NoMemory();
+        return -1;
+    }
+    /* The first item of vectorcall_args will be filled with obj later */
+    if (nargs > 1) {
+        memcpy(mc->vectorcall_args, PySequence_Fast_ITEMS(args),
+            nargs * sizeof(PyObject*));
+    }
+    if (kwds) {
+        const Py_ssize_t nkwds = PyDict_Size(kwds);
+
+        mc->vectorcall_kwnames = PyTuple_New(nkwds);
+        if (!mc->vectorcall_kwnames) {
+            return -1;
+        }
+        Py_ssize_t i = 0, ppos = 0;
+        PyObject* key, * value;
+        while (PyDict_Next(kwds, &ppos, &key, &value)) {
+            PyTuple_SET_ITEM(mc->vectorcall_kwnames, i, Py_NewRef(key));
+            mc->vectorcall_args[nargs + i] = value; // borrowed reference
+            ++i;
+        }
+    }
+    else {
+        mc->vectorcall_kwnames = NULL;
+    }
+    return 1;
+}
+
+
+static PyObject *
+methodcaller_vectorcall(
+        methodcallerobject *mc, PyObject *const *args, size_t nargsf, PyObject* kwnames)
+{
+    if (!_PyArg_CheckPositional("methodcaller", PyVectorcall_NARGS(nargsf), 1, 1)
+        || !_PyArg_NoKwnames("methodcaller", kwnames)) {
+        return NULL;
+    }
+    if (mc->vectorcall_args == NULL) {
+        if (_methodcaller_initialize_vectorcall(mc) < 0) {
+            return NULL;
+        }
+    }
+
+    assert(mc->vectorcall_args != 0);
+    mc->vectorcall_args[0] = args[0];
+    return PyObject_VectorcallMethod(
+            mc->name, mc->vectorcall_args,
+            (PyTuple_GET_SIZE(mc->xargs)) | PY_VECTORCALL_ARGUMENTS_OFFSET,
+            mc->vectorcall_kwnames);
+}
+
+
 /* AC 3.5: variable number of arguments, not currently support by AC */
 static PyObject *
 methodcaller_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
@@ -1580,30 +1647,32 @@ methodcaller_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
         return NULL;
     }
 
-    name = PyTuple_GET_ITEM(args, 0);
     Py_INCREF(name);
     PyUnicode_InternInPlace(&name);
     mc->name = name;
 
+    mc->xargs = Py_XNewRef(args); // allows us to use borrowed references
     mc->kwds = Py_XNewRef(kwds);
+    mc->vectorcall_args = 0;
 
-    mc->args = PyTuple_GetSlice(args, 1, PyTuple_GET_SIZE(args));
-    if (mc->args == NULL) {
-        Py_DECREF(mc);
-        return NULL;
-    }
+
+    mc->vectorcall = (vectorcallfunc)methodcaller_vectorcall;
 
     PyObject_GC_Track(mc);
     return (PyObject *)mc;
 }
 
-static int
+static void
 methodcaller_clear(methodcallerobject *mc)
 {
     Py_CLEAR(mc->name);
-    Py_CLEAR(mc->args);
+    Py_CLEAR(mc->xargs);
     Py_CLEAR(mc->kwds);
-    return 0;
+    if (mc->vectorcall_args != NULL) {
+        PyMem_Free(mc->vectorcall_args);
+        mc->vectorcall_args = 0;
+        Py_CLEAR(mc->vectorcall_kwnames);
+    }
 }
 
 static void
@@ -1611,7 +1680,7 @@ methodcaller_dealloc(methodcallerobject *mc)
 {
     PyTypeObject *tp = Py_TYPE(mc);
     PyObject_GC_UnTrack(mc);
-    (void)methodcaller_clear(mc);
+    methodcaller_clear(mc);
     tp->tp_free(mc);
     Py_DECREF(tp);
 }
@@ -1620,7 +1689,7 @@ static int
 methodcaller_traverse(methodcallerobject *mc, visitproc visit, void *arg)
 {
     Py_VISIT(mc->name);
-    Py_VISIT(mc->args);
+    Py_VISIT(mc->xargs);
     Py_VISIT(mc->kwds);
     Py_VISIT(Py_TYPE(mc));
     return 0;
@@ -1639,7 +1708,16 @@ methodcaller_call(methodcallerobject *mc, PyObject *args, PyObject *kw)
     method = PyObject_GetAttr(obj, mc->name);
     if (method == NULL)
         return NULL;
-    result = PyObject_Call(method, mc->args, mc->kwds);
+
+
+    PyObject *cargs = PyTuple_GetSlice(mc->xargs, 1, PyTuple_GET_SIZE(mc->xargs));
+    if (cargs == NULL) {
+        Py_DECREF(method);
+        return NULL;
+    }
+
+    result = PyObject_Call(method, cargs, mc->kwds);
+    Py_DECREF(cargs);
     Py_DECREF(method);
     return result;
 }
@@ -1657,7 +1735,7 @@ methodcaller_repr(methodcallerobject *mc)
     }
 
     numkwdargs = mc->kwds != NULL ? PyDict_GET_SIZE(mc->kwds) : 0;
-    numposargs = PyTuple_GET_SIZE(mc->args);
+    numposargs = PyTuple_GET_SIZE(mc->xargs) - 1;
     numtotalargs = numposargs + numkwdargs;
 
     if (numtotalargs == 0) {
@@ -1673,7 +1751,7 @@ methodcaller_repr(methodcallerobject *mc)
     }
 
     for (i = 0; i < numposargs; ++i) {
-        PyObject *onerepr = PyObject_Repr(PyTuple_GET_ITEM(mc->args, i));
+        PyObject *onerepr = PyObject_Repr(PyTuple_GET_ITEM(mc->xargs, i+1));
         if (onerepr == NULL)
             goto done;
         PyTuple_SET_ITEM(argreprs, i, onerepr);
@@ -1723,17 +1801,16 @@ methodcaller_repr(methodcallerobject *mc)
 static PyObject *
 methodcaller_reduce(methodcallerobject *mc, PyObject *Py_UNUSED(ignored))
 {
-    PyObject *newargs;
     if (!mc->kwds || PyDict_GET_SIZE(mc->kwds) == 0) {
         Py_ssize_t i;
-        Py_ssize_t callargcount = PyTuple_GET_SIZE(mc->args);
-        newargs = PyTuple_New(1 + callargcount);
+        Py_ssize_t newarg_size = PyTuple_GET_SIZE(mc->xargs);
+        PyObject *newargs = PyTuple_New(newarg_size);
         if (newargs == NULL)
             return NULL;
         PyTuple_SET_ITEM(newargs, 0, Py_NewRef(mc->name));
-        for (i = 0; i < callargcount; ++i) {
-            PyObject *arg = PyTuple_GET_ITEM(mc->args, i);
-            PyTuple_SET_ITEM(newargs, i + 1, Py_NewRef(arg));
+        for (i = 1; i < newarg_size; ++i) {
+            PyObject *arg = PyTuple_GET_ITEM(mc->xargs, i);
+            PyTuple_SET_ITEM(newargs, i, Py_NewRef(arg));
         }
         return Py_BuildValue("ON", Py_TYPE(mc), newargs);
     }
@@ -1751,7 +1828,12 @@ methodcaller_reduce(methodcallerobject *mc, PyObject *Py_UNUSED(ignored))
         constructor = PyObject_VectorcallDict(partial, newargs, 2, mc->kwds);
 
         Py_DECREF(partial);
-        return Py_BuildValue("NO", constructor, mc->args);
+        PyObject *args = PyTuple_GetSlice(mc->xargs, 1, PyTuple_GET_SIZE(mc->xargs));
+        if (!args) {
+            Py_DECREF(constructor);
+            return NULL;
+        }
+        return Py_BuildValue("NO", constructor, args);
     }
 }
 
@@ -1760,6 +1842,12 @@ static PyMethodDef methodcaller_methods[] = {
      reduce_doc},
     {NULL}
 };
+
+static PyMemberDef methodcaller_members[] = {
+    {"__vectorcalloffset__", Py_T_PYSSIZET, offsetof(methodcallerobject, vectorcall), Py_READONLY},
+    {NULL}
+};
+
 PyDoc_STRVAR(methodcaller_doc,
 "methodcaller(name, /, *args, **kwargs)\n--\n\n\
 Return a callable object that calls the given method on its operand.\n\
@@ -1774,6 +1862,7 @@ static PyType_Slot methodcaller_type_slots[] = {
     {Py_tp_traverse, methodcaller_traverse},
     {Py_tp_clear, methodcaller_clear},
     {Py_tp_methods, methodcaller_methods},
+    {Py_tp_members, methodcaller_members},
     {Py_tp_new, methodcaller_new},
     {Py_tp_getattro, PyObject_GenericGetAttr},
     {Py_tp_repr, methodcaller_repr},
@@ -1785,7 +1874,7 @@ static PyType_Spec methodcaller_type_spec = {
     .basicsize = sizeof(methodcallerobject),
     .itemsize = 0,
     .flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
-              Py_TPFLAGS_IMMUTABLETYPE),
+              Py_TPFLAGS_HAVE_VECTORCALL | Py_TPFLAGS_IMMUTABLETYPE),
     .slots = methodcaller_type_slots,
 };