Sep 26, 2024 · Sep 26, 2024 · Sep 26, 2024 · Sep 26, 2024 · Sep 27, 2024 · Sep 27, 2024
diff --git a/Misc/NEWS.d/next/Library/2024-10-17-01-12-22.gh-issue-119109.u4hcvb.rst b/Misc/NEWS.d/next/Library/2024-10-17-01-12-22.gh-issue-119109.u4hcvb.rst
 :func:`functools.partial` calls are now faster when keyword arguments are used.
diff --git a/Modules/_functoolsmodule.c b/Modules/_functoolsmodule.c
    return PyMethod_New(self, obj);
 }

 /* Merging keyword arguments using the vectorcall convention is messy, so
 * if we would need to do that, we stop using vectorcall and fall back
 * to using partial_call() instead. */
 Py_NO_INLINE static PyObject *
 partial_vectorcall_fallback(PyThreadState *tstate, partialobject *pto,
                            PyObject *const *args, size_t nargsf,
                            PyObject *kwnames)
 {
    pto->vectorcall = NULL;
    Py_ssize_t nargs = PyVectorcall_NARGS(nargsf);
    return _PyObject_MakeTpCall(tstate, (PyObject *)pto, args, nargs, kwnames);
 }

 static PyObject *
 partial_vectorcall(PyObject *self, PyObject *const *args,
                   size_t nargsf, PyObject *kwnames)
 {
    partialobject *pto = _PyPartialObject_CAST(self);;
    PyThreadState *tstate = _PyThreadState_GET();
    /* Sizes */
    Py_ssize_t pto_nargs = PyTuple_GET_SIZE(pto->args);
    Py_ssize_t pto_nkwds = PyDict_GET_SIZE(pto->kw);
    Py_ssize_t nargs = PyVectorcall_NARGS(nargsf);
    Py_ssize_t nkwds = kwnames == NULL ? 0 : PyTuple_GET_SIZE(kwnames);
    Py_ssize_t nargskw = nargs + nkwds;

    /* pto->kw is mutable, so need to check every time */
    if (PyDict_GET_SIZE(pto->kw)) {
        return partial_vectorcall_fallback(tstate, pto, args, nargsf, kwnames);
    }
    /* Placeholder check */
    Py_ssize_t pto_phcount = pto->phcount;
    if (nargs < pto_phcount) {
        PyErr_Format(PyExc_TypeError,
        return NULL;
    }

    Py_ssize_t nargskw = nargs;
    if (kwnames != NULL) {
        nargskw += PyTuple_GET_SIZE(kwnames);
    }

    PyObject **pto_args = _PyTuple_ITEMS(pto->args);
    Py_ssize_t pto_nargs = PyTuple_GET_SIZE(pto->args);

    /* Fast path if we're called without arguments */
    if (nargskw == 0) {
        return _PyObject_VectorcallTstate(tstate, pto->fn,
                                          pto_args, pto_nargs, NULL);
    }
    /* Special cases */
    if (!pto_nkwds) {
        /* Fast path if we're called without arguments */
        if (nargskw == 0) {
            return _PyObject_VectorcallTstate(tstate, pto->fn, pto_args,
                                              pto_nargs, NULL);
        }

    /* Fast path using PY_VECTORCALL_ARGUMENTS_OFFSET to prepend a single
     * positional argument */
    if (pto_nargs == 1 && (nargsf & PY_VECTORCALL_ARGUMENTS_OFFSET)) {
        PyObject **newargs = (PyObject **)args - 1;
        PyObject *tmp = newargs[0];
        newargs[0] = pto_args[0];
        PyObject *ret = _PyObject_VectorcallTstate(tstate, pto->fn,
                                                   newargs, nargs + 1, kwnames);
        newargs[0] = tmp;
        return ret;
        // TO BE ADDED ON MERGE TO MAIN
        //      IF https://github.com/python/cpython/pull/124788
        //      IS MERGED BEFORE THIS
        // /* Fast path if all Placeholders */
        // if (pto_nargs == pto_phcount) {
        //     /* NOTE: Without this, following single argument Fast Path
        //      * is incorrect */
        //     return _PyObject_VectorcallTstate(tstate, pto->fn,
        //                                       args, nargs, kwnames);
        // }

        /* Use PY_VECTORCALL_ARGUMENTS_OFFSET to prepend a single
         * positional argument. */
        if (pto_nargs == 1 && (nargsf & PY_VECTORCALL_ARGUMENTS_OFFSET)) {
            PyObject **newargs = (PyObject **)args - 1;
            PyObject *tmp = newargs[0];
            newargs[0] = pto_args[0];
            PyObject *ret = _PyObject_VectorcallTstate(tstate, pto->fn, newargs,
                                                       nargs + 1, kwnames);
            newargs[0] = tmp;
            return ret;
        }
    }

    PyObject *small_stack[_PY_FASTCALL_SMALL_STACK];
    PyObject **stack;
    /* Total sizes */
    Py_ssize_t tot_nargs = pto_nargs + nargs - pto_phcount;
    Py_ssize_t tot_nkwds = pto_nkwds + nkwds;
    Py_ssize_t tot_nargskw = tot_nargs + tot_nkwds;

    PyObject *pto_kw_merged = NULL;  // pto_kw with duplicates merged (if any)
    PyObject *tot_kwnames;

    Py_ssize_t tot_nargskw = pto_nargs + nargskw - pto_phcount;
    if (tot_nargskw <= (Py_ssize_t)Py_ARRAY_LENGTH(small_stack)) {
    /* Allocate Stack */
    PyObject **tmp_stack, **stack, *small_stack[_PY_FASTCALL_SMALL_STACK];
    Py_ssize_t init_stack_size = tot_nargskw;
    if (pto_nkwds) {
        // If pto_nkwds, allocate additional space for temporary new keys
        init_stack_size += nkwds;
    }
    if (init_stack_size <= (Py_ssize_t)Py_ARRAY_LENGTH(small_stack)) {
        stack = small_stack;
    }
    else {
        stack = PyMem_Malloc(tot_nargskw * sizeof(PyObject *));
        /* NOTE, in theory size * elsize could overflow */
        stack = PyMem_Malloc(init_stack_size * sizeof(PyObject *));
        if (stack == NULL) {
            PyErr_NoMemory();
            return NULL;
        }
    }

    Py_ssize_t tot_nargs;
    /* Copy keywords to stack */
    if (!pto_nkwds) {
        tot_kwnames = kwnames;
        if (nkwds) {
            /* if !pto_nkwds & nkwds, then simply append kw */
            memcpy(stack + tot_nargs, args + nargs, nkwds * sizeof(PyObject*));
        }
    }
    else {
        /* stack is now     [<positionals>, <pto_kwds>, <kwds>, <kwds_keys>]
         * Will truncate to [<positionals>, <merged_kwds>] */
        PyObject *key, *val;

        /* Merge kw to pto_kw or add to tail (if not duplicate) */
        Py_ssize_t n_tail = 0;
        for (Py_ssize_t i = 0; i < nkwds; ++i) {
            key = PyTuple_GET_ITEM(kwnames, i);
            val = args[nargs + i];
            if (PyDict_Contains(pto->kw, key)) {
                if (pto_kw_merged == NULL) {
                    pto_kw_merged = PyDict_Copy(pto->kw);
                    if (pto_kw_merged == NULL) {
                        goto error_2;
                    }
                }
                if (PyDict_SetItem(pto_kw_merged, key, val) < 0) {
                    goto error_1;
                }
            }
            else {
                /* Copy keyword tail to stack */
                stack[tot_nargs + pto_nkwds + n_tail] = val;
                stack[tot_nargskw + n_tail] = key;
                n_tail++;
            }
        }
        Py_ssize_t n_merges = nkwds - n_tail;

        /* Create tot_kwnames */
        tot_kwnames = PyTuple_New(pto_nkwds + n_tail);
        for (Py_ssize_t i = 0; i < n_tail; ++i) {
            key = stack[tot_nargskw + i];
            PyTuple_SET_ITEM(tot_kwnames, pto_nkwds + i, key);
            Py_INCREF(key);
        }

        /* Copy pto_kw_merged to stack */
        Py_ssize_t pos = 0, i = 0;
        while (PyDict_Next(n_merges ? pto_kw_merged : pto->kw, &pos, &key, &val)) {
            PyTuple_SET_ITEM(tot_kwnames, i, key);
            Py_INCREF(key);
            stack[tot_nargs + i] = val;
            i++;
        }

        /* Resize Stack */
        if (n_merges && stack != small_stack) {
            tmp_stack = PyMem_Realloc(stack, (tot_nargskw - n_merges) * sizeof(PyObject *));
            if (tmp_stack == NULL) {
                PyErr_NoMemory();
                goto error_0;
            }
            stack = tmp_stack;
        }

        Py_XDECREF(pto_kw_merged);
    }

    /* Copy Positionals to stack */
    if (pto_phcount) {
        tot_nargs = pto_nargs + nargs - pto_phcount;
        Py_ssize_t j = 0;       // New args index
        for (Py_ssize_t i = 0; i < pto_nargs; i++) {
            if (pto_args[i] == pto->placeholder){
            if (pto_args[i] == pto->placeholder){
                stack[i] = args[j];
                j += 1;
            }
            }
        }
        assert(j == pto_phcount);
        if (nargskw > pto_phcount) {
            memcpy(stack + pto_nargs, args + j, (nargskw - j) * sizeof(PyObject*));
        /* Add remaining args from new_args */
        if (nargs > pto_phcount) {
            memcpy(stack + pto_nargs, args + j, (nargs - j) * sizeof(PyObject*));
        }
    }
    else {
        tot_nargs = pto_nargs + nargs;
        /* Copy to new stack, using borrowed references */
        memcpy(stack, pto_args, pto_nargs * sizeof(PyObject*));
        memcpy(stack + pto_nargs, args,nargskw * sizeof(PyObject*));
        memcpy(stack + pto_nargs, args,nargs * sizeof(PyObject*));
    }
    PyObject *ret = _PyObject_VectorcallTstate(tstate, pto->fn,
                                               stack, tot_nargs, kwnames);

    /* Call / Maintenance / Return */
    PyObject *ret = _PyObject_VectorcallTstate(tstate, pto->fn, stack,
                                               tot_nargs, tot_kwnames);
    if (stack != small_stack) {
        PyMem_Free(stack);
    }
    if (pto_nkwds) {
        Py_DECREF(tot_kwnames);
    }
    return ret;

 error_0:
    Py_DECREF(tot_kwnames);
 error_1:
    Py_XDECREF(pto_kw_merged);
 error_2:
    if (stack != small_stack) {
        PyMem_Free(stack);
    }
    return NULL;
 }

 /* Set pto->vectorcall depending on the parameters of the partial object */
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		:func:`functools.partial` calls are now faster when keyword arguments are used.
dg-pb marked this conversation as resolved. Show resolvedHide resolved
Original file line number	Diff line number	Diff line change
Expand Up		@@ -355,31 +355,20 @@ partial_descr_get(PyObject self, PyObject obj, PyObject *type)
		return PyMethod_New(self, obj);
		}

		/* Merging keyword arguments using the vectorcall convention is messy, so
		* if we would need to do that, we stop using vectorcall and fall back
		* to using partial_call() instead. */
		Py_NO_INLINE static PyObject *
		partial_vectorcall_fallback(PyThreadState tstate, partialobject pto,
		PyObject const args, size_t nargsf,
		PyObject *kwnames)
		{
		pto->vectorcall = NULL;
		Py_ssize_t nargs = PyVectorcall_NARGS(nargsf);
		return _PyObject_MakeTpCall(tstate, (PyObject *)pto, args, nargs, kwnames);
		}

		static PyObject *
		partial_vectorcall(PyObject self, PyObject const *args,
		size_t nargsf, PyObject *kwnames)
		{
		partialobject *pto = _PyPartialObject_CAST(self);;
		PyThreadState *tstate = _PyThreadState_GET();
		/* Sizes */
dg-pb marked this conversation as resolved. Show resolvedHide resolved
		Py_ssize_t pto_nargs = PyTuple_GET_SIZE(pto->args);
		Py_ssize_t pto_nkwds = PyDict_GET_SIZE(pto->kw);
		Py_ssize_t nargs = PyVectorcall_NARGS(nargsf);
		Py_ssize_t nkwds = kwnames == NULL ? 0 : PyTuple_GET_SIZE(kwnames);
		Py_ssize_t nargskw = nargs + nkwds;

		/* pto->kw is mutable, so need to check every time */
		if (PyDict_GET_SIZE(pto->kw)) {
		return partial_vectorcall_fallback(tstate, pto, args, nargsf, kwnames);
		}
		/* Placeholder check */
		Py_ssize_t pto_phcount = pto->phcount;
		if (nargs < pto_phcount) {
		PyErr_Format(PyExc_TypeError,
Expand All		@@ -388,53 +377,140 @@ partial_vectorcall(PyObject self, PyObject const *args,
		return NULL;
		}

		Py_ssize_t nargskw = nargs;
		if (kwnames != NULL) {
		nargskw += PyTuple_GET_SIZE(kwnames);
		}

		PyObject **pto_args = _PyTuple_ITEMS(pto->args);
		Py_ssize_t pto_nargs = PyTuple_GET_SIZE(pto->args);

		/* Fast path if we're called without arguments */
		if (nargskw == 0) {
		return _PyObject_VectorcallTstate(tstate, pto->fn,
		pto_args, pto_nargs, NULL);
		}
		/* Special cases */
		if (!pto_nkwds) {
		/* Fast path if we're called without arguments */
		if (nargskw == 0) {
		return _PyObject_VectorcallTstate(tstate, pto->fn, pto_args,
		pto_nargs, NULL);
		}

		/* Fast path using PY_VECTORCALL_ARGUMENTS_OFFSET to prepend a single
		* positional argument */
		if (pto_nargs == 1 && (nargsf & PY_VECTORCALL_ARGUMENTS_OFFSET)) {
		PyObject newargs = (PyObject )args - 1;
		PyObject *tmp = newargs[0];
		newargs[0] = pto_args[0];
		PyObject *ret = _PyObject_VectorcallTstate(tstate, pto->fn,
		newargs, nargs + 1, kwnames);
		newargs[0] = tmp;
		return ret;
		// TO BE ADDED ON MERGE TO MAIN
		// IF https://github.com/python/cpython/pull/124788
		// IS MERGED BEFORE THIS
		// /* Fast path if all Placeholders */
		// if (pto_nargs == pto_phcount) {
		// /* NOTE: Without this, following single argument Fast Path
		// * is incorrect */
		// return _PyObject_VectorcallTstate(tstate, pto->fn,
		// args, nargs, kwnames);
		// }

		/* Use PY_VECTORCALL_ARGUMENTS_OFFSET to prepend a single
		* positional argument. */
		if (pto_nargs == 1 && (nargsf & PY_VECTORCALL_ARGUMENTS_OFFSET)) {
		PyObject newargs = (PyObject )args - 1;
		PyObject *tmp = newargs[0];
		newargs[0] = pto_args[0];
		PyObject *ret = _PyObject_VectorcallTstate(tstate, pto->fn, newargs,
		nargs + 1, kwnames);
		newargs[0] = tmp;
		return ret;
		}
		}

		PyObject *small_stack[_PY_FASTCALL_SMALL_STACK];
		PyObject **stack;
		/* Total sizes */
		Py_ssize_t tot_nargs = pto_nargs + nargs - pto_phcount;
		Py_ssize_t tot_nkwds = pto_nkwds + nkwds;
		Py_ssize_t tot_nargskw = tot_nargs + tot_nkwds;

		PyObject *pto_kw_merged = NULL; // pto_kw with duplicates merged (if any)
		PyObject *tot_kwnames;

		Py_ssize_t tot_nargskw = pto_nargs + nargskw - pto_phcount;
		if (tot_nargskw <= (Py_ssize_t)Py_ARRAY_LENGTH(small_stack)) {
		/* Allocate Stack */
		PyObject tmp_stack, stack, *small_stack[_PY_FASTCALL_SMALL_STACK];
		Py_ssize_t init_stack_size = tot_nargskw;
		if (pto_nkwds) {
		// If pto_nkwds, allocate additional space for temporary new keys
		init_stack_size += nkwds;
		}
		if (init_stack_size <= (Py_ssize_t)Py_ARRAY_LENGTH(small_stack)) {
		stack = small_stack;
		}
		else {
		stack = PyMem_Malloc(tot_nargskw * sizeof(PyObject *));
		/* NOTE, in theory size * elsize could overflow */
		stack = PyMem_Malloc(init_stack_size * sizeof(PyObject *));
		if (stack == NULL) {
		PyErr_NoMemory();
		return NULL;
		}
		}

		Py_ssize_t tot_nargs;
		/* Copy keywords to stack */
		if (!pto_nkwds) {
		tot_kwnames = kwnames;
		if (nkwds) {
		/* if !pto_nkwds & nkwds, then simply append kw */
		memcpy(stack + tot_nargs, args + nargs, nkwds * sizeof(PyObject*));
dg-pb marked this conversation as resolved. Show resolvedHide resolved
		}
		}
		else {
		/* stack is now [<positionals>, <pto_kwds>, <kwds>, <kwds_keys>]
		* Will truncate to [<positionals>, <merged_kwds>] */
		PyObject key, val;

		/* Merge kw to pto_kw or add to tail (if not duplicate) */
		Py_ssize_t n_tail = 0;
		for (Py_ssize_t i = 0; i < nkwds; ++i) {
		key = PyTuple_GET_ITEM(kwnames, i);
		val = args[nargs + i];
		if (PyDict_Contains(pto->kw, key)) {
		if (pto_kw_merged == NULL) {
		pto_kw_merged = PyDict_Copy(pto->kw);
		if (pto_kw_merged == NULL) {
		goto error_2;
		}
		}
		if (PyDict_SetItem(pto_kw_merged, key, val) < 0) {
		goto error_1;
		}
		}
		else {
		/* Copy keyword tail to stack */
		stack[tot_nargs + pto_nkwds + n_tail] = val;
		stack[tot_nargskw + n_tail] = key;
		n_tail++;
		}
		}
		Py_ssize_t n_merges = nkwds - n_tail;

		/* Create tot_kwnames */
		tot_kwnames = PyTuple_New(pto_nkwds + n_tail);
		for (Py_ssize_t i = 0; i < n_tail; ++i) {
		key = stack[tot_nargskw + i];
		PyTuple_SET_ITEM(tot_kwnames, pto_nkwds + i, key);
		Py_INCREF(key);
dg-pb marked this conversation as resolved. Show resolvedHide resolved
		}

		/* Copy pto_kw_merged to stack */
		Py_ssize_t pos = 0, i = 0;
		while (PyDict_Next(n_merges ? pto_kw_merged : pto->kw, &pos, &key, &val)) {
		PyTuple_SET_ITEM(tot_kwnames, i, key);
		Py_INCREF(key);
dg-pb marked this conversation as resolved. Show resolvedHide resolved
		stack[tot_nargs + i] = val;
		i++;
		}

		/* Resize Stack */
		if (n_merges && stack != small_stack) {
		tmp_stack = PyMem_Realloc(stack, (tot_nargskw - n_merges) * sizeof(PyObject *));
		if (tmp_stack == NULL) {
		PyErr_NoMemory();
		goto error_0;
		}
		stack = tmp_stack;
		}

		Py_XDECREF(pto_kw_merged);
		}

		/* Copy Positionals to stack */
		if (pto_phcount) {
		tot_nargs = pto_nargs + nargs - pto_phcount;
		Py_ssize_t j = 0; // New args index
		for (Py_ssize_t i = 0; i < pto_nargs; i++) {
		if (pto_args[i] == pto->placeholder){
		if (pto_args[i] == pto->placeholder){
		stack[i] = args[j];
		j += 1;
		}
Expand All		@@ -443,22 +519,36 @@ partial_vectorcall(PyObject self, PyObject const *args,
		}
		}
		assert(j == pto_phcount);
		if (nargskw > pto_phcount) {
		memcpy(stack + pto_nargs, args + j, (nargskw - j) * sizeof(PyObject*));
		/* Add remaining args from new_args */
		if (nargs > pto_phcount) {
		memcpy(stack + pto_nargs, args + j, (nargs - j) * sizeof(PyObject*));
dg-pb marked this conversation as resolved. Show resolvedHide resolved
		}
		}
		else {
		tot_nargs = pto_nargs + nargs;
		/* Copy to new stack, using borrowed references */
		memcpy(stack, pto_args, pto_nargs * sizeof(PyObject*));
		memcpy(stack + pto_nargs, args,nargskw * sizeof(PyObject*));
		memcpy(stack + pto_nargs, args,nargs * sizeof(PyObject*));
dg-pb marked this conversation as resolved. Show resolvedHide resolved
		}
		PyObject *ret = _PyObject_VectorcallTstate(tstate, pto->fn,
		stack, tot_nargs, kwnames);

		/* Call / Maintenance / Return */
		PyObject *ret = _PyObject_VectorcallTstate(tstate, pto->fn, stack,
		tot_nargs, tot_kwnames);
		if (stack != small_stack) {
		PyMem_Free(stack);
		}
		if (pto_nkwds) {
		Py_DECREF(tot_kwnames);
		}
		return ret;

		error_0:
		Py_DECREF(tot_kwnames);
		error_1:
		Py_XDECREF(pto_kw_merged);
		error_2:
		if (stack != small_stack) {
		PyMem_Free(stack);
		}
erlend-aasland marked this conversation as resolved. Show resolvedHide resolved
		return NULL;
		}

		/* Set pto->vectorcall depending on the parameters of the partial object */
Expand Down