bpo-38644: Add _PyObject_VectorcallTstate() (GH-17052)

* Add _PyObject_VectorcallTstate() function: similar to
  _PyObject_Vectorcall(), but with tstate parameter
* Add tstate parameter to _PyObject_MakeTpCall()

6 files changed, 94 insertions, 50 deletions

diff --git a/Include/cpython/abstract.h b/Include/cpython/abstract.h
index be37e19..fef538e 100644
--- a/Include/cpython/abstract.h
+++ b/Include/cpython/abstract.h
@@ -49,6 +49,7 @@ PyAPI_FUNC(PyObject *) _Py_CheckFunctionResult(
    or _PyObject_FastCallDict() (both forms are supported),
    except that nargs is plainly the number of arguments without flags. */
 PyAPI_FUNC(PyObject *) _PyObject_MakeTpCall(
+    PyThreadState *tstate,
     PyObject *callable,
     PyObject *const *args, Py_ssize_t nargs,
     PyObject *keywords);
@@ -95,22 +96,31 @@ _PyVectorcall_Function(PyObject *callable)
    Return the result on success. Raise an exception and return NULL on
    error. */
 static inline PyObject *
-_PyObject_Vectorcall(PyObject *callable, PyObject *const *args,
-                     size_t nargsf, PyObject *kwnames)
+_PyObject_VectorcallTstate(PyThreadState *tstate, PyObject *callable,
+                           PyObject *const *args, size_t nargsf,
+                           PyObject *kwnames)
 {
     assert(kwnames == NULL || PyTuple_Check(kwnames));
     assert(args != NULL || PyVectorcall_NARGS(nargsf) == 0);
 
-    PyThreadState *tstate = PyThreadState_GET();
     vectorcallfunc func = _PyVectorcall_Function(callable);
     if (func == NULL) {
         Py_ssize_t nargs = PyVectorcall_NARGS(nargsf);
-        return _PyObject_MakeTpCall(callable, args, nargs, kwnames);
+        return _PyObject_MakeTpCall(tstate, callable, args, nargs, kwnames);
     }
     PyObject *res = func(callable, args, nargsf, kwnames);
     return _Py_CheckFunctionResult(tstate, callable, res, NULL);
 }
 
+static inline PyObject *
+_PyObject_Vectorcall(PyObject *callable, PyObject *const *args,
+                     size_t nargsf, PyObject *kwnames)
+{
+    PyThreadState *tstate = PyThreadState_GET();
+    return _PyObject_VectorcallTstate(tstate, callable,
+                                      args, nargsf, kwnames);
+}
+
 /* Same as _PyObject_Vectorcall except that keyword arguments are passed as
    dict, which may be NULL if there are no keyword arguments. */
 PyAPI_FUNC(PyObject *) _PyObject_FastCallDict(
diff --git a/Modules/_functoolsmodule.c b/Modules/_functoolsmodule.c
index 17d49ac..987087b 100644
--- a/Modules/_functoolsmodule.c
+++ b/Modules/_functoolsmodule.c
@@ -132,21 +132,25 @@ partial_dealloc(partialobject *pto)
  * 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(partialobject *pto, PyObject *const *args,
-                            size_t nargsf, PyObject *kwnames)
+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((PyObject *)pto, args, nargs, kwnames);
+    return _PyObject_MakeTpCall(tstate, (PyObject *)pto,
+                                args, nargs, kwnames);
 }
 
 static PyObject *
 partial_vectorcall(partialobject *pto, PyObject *const *args,
                    size_t nargsf, PyObject *kwnames)
 {
+    PyThreadState *tstate = _PyThreadState_GET();
+
     /* pto->kw is mutable, so need to check every time */
     if (PyDict_GET_SIZE(pto->kw)) {
-        return partial_vectorcall_fallback(pto, args, nargsf, kwnames);
+        return partial_vectorcall_fallback(tstate, pto, args, nargsf, kwnames);
     }
 
     Py_ssize_t nargs = PyVectorcall_NARGS(nargsf);
@@ -160,7 +164,8 @@ partial_vectorcall(partialobject *pto, PyObject *const *args,
 
     /* Fast path if we're called without arguments */
     if (nargs_total == 0) {
-        return _PyObject_Vectorcall(pto->fn, pto_args, pto_nargs, NULL);
+        return _PyObject_VectorcallTstate(tstate, pto->fn,
+                                          pto_args, pto_nargs, NULL);
     }
 
     /* Fast path using PY_VECTORCALL_ARGUMENTS_OFFSET to prepend a single
@@ -169,7 +174,8 @@ partial_vectorcall(partialobject *pto, PyObject *const *args,
         PyObject **newargs = (PyObject **)args - 1;
         PyObject *tmp = newargs[0];
         newargs[0] = pto_args[0];
-        PyObject *ret = _PyObject_Vectorcall(pto->fn, newargs, nargs + 1, kwnames);
+        PyObject *ret = _PyObject_VectorcallTstate(tstate, pto->fn,
+                                                   newargs, nargs + 1, kwnames);
         newargs[0] = tmp;
         return ret;
     }
@@ -195,7 +201,8 @@ partial_vectorcall(partialobject *pto, PyObject *const *args,
     memcpy(stack, pto_args, pto_nargs * sizeof(PyObject*));
     memcpy(stack + pto_nargs, args, nargs_total * sizeof(PyObject*));
 
-    ret = _PyObject_Vectorcall(pto->fn, stack, pto_nargs + nargs, kwnames);
+    ret = _PyObject_VectorcallTstate(tstate, pto->fn,
+                                     stack, pto_nargs + nargs, kwnames);
     if (stack != small_stack) {
         PyMem_Free(stack);
     }
diff --git a/Objects/call.c b/Objects/call.c
index a1d0b33..a8ae41a 100644
--- a/Objects/call.c
+++ b/Objects/call.c
@@ -104,7 +104,7 @@ _PyObject_FastCallDict(PyObject *callable, PyObject *const *args,
     vectorcallfunc func = _PyVectorcall_Function(callable);
     if (func == NULL) {
         /* Use tp_call instead */
-        return _PyObject_MakeTpCall(callable, args, nargs, kwargs);
+        return _PyObject_MakeTpCall(tstate, callable, args, nargs, kwargs);
     }
 
     PyObject *res;
@@ -129,10 +129,10 @@ _PyObject_FastCallDict(PyObject *callable, PyObject *const *args,
 
 
 PyObject *
-_PyObject_MakeTpCall(PyObject *callable, PyObject *const *args, Py_ssize_t nargs, PyObject *keywords)
+_PyObject_MakeTpCall(PyThreadState *tstate, PyObject *callable,
+                     PyObject *const *args, Py_ssize_t nargs,
+                     PyObject *keywords)
 {
-    PyThreadState *tstate = _PyThreadState_GET();
-
     /* Slow path: build a temporary tuple for positional arguments and a
      * temporary dictionary for keyword arguments (if any) */
     ternaryfunc call = Py_TYPE(callable)->tp_call;
@@ -774,6 +774,7 @@ _PyObject_VectorcallMethod(PyObject *name, PyObject *const *args,
     assert(args != NULL);
     assert(PyVectorcall_NARGS(nargsf) >= 1);
 
+    PyThreadState *tstate = _PyThreadState_GET();
     PyObject *callable = NULL;
     /* Use args[0] as "self" argument */
     int unbound = _PyObject_GetMethod(args[0], name, &callable);
@@ -792,7 +793,8 @@ _PyObject_VectorcallMethod(PyObject *name, PyObject *const *args,
         args++;
         nargsf--;
     }
-    PyObject *result = _PyObject_Vectorcall(callable, args, nargsf, kwnames);
+    PyObject *result = _PyObject_VectorcallTstate(tstate, callable,
+                                                  args, nargsf, kwnames);
     Py_DECREF(callable);
     return result;
 }
diff --git a/Objects/classobject.c b/Objects/classobject.c
index 4a9add1..d3fc726 100644
--- a/Objects/classobject.c
+++ b/Objects/classobject.c
@@ -2,6 +2,7 @@
 
 #include "Python.h"
 #include "pycore_object.h"
+#include "pycore_pyerrors.h"
 #include "pycore_pymem.h"
 #include "pycore_pystate.h"
 #include "structmember.h"
@@ -37,25 +38,28 @@ method_vectorcall(PyObject *method, PyObject *const *args,
                   size_t nargsf, PyObject *kwnames)
 {
     assert(Py_TYPE(method) == &PyMethod_Type);
-    PyObject *self, *func, *result;
-    self = PyMethod_GET_SELF(method);
-    func = PyMethod_GET_FUNCTION(method);
+
+    PyThreadState *tstate = _PyThreadState_GET();
+    PyObject *self = PyMethod_GET_SELF(method);
+    PyObject *func = PyMethod_GET_FUNCTION(method);
     Py_ssize_t nargs = PyVectorcall_NARGS(nargsf);
 
+    PyObject *result;
     if (nargsf & PY_VECTORCALL_ARGUMENTS_OFFSET) {
         /* PY_VECTORCALL_ARGUMENTS_OFFSET is set, so we are allowed to mutate the vector */
         PyObject **newargs = (PyObject**)args - 1;
         nargs += 1;
         PyObject *tmp = newargs[0];
         newargs[0] = self;
-        result = _PyObject_Vectorcall(func, newargs, nargs, kwnames);
+        result = _PyObject_VectorcallTstate(tstate, func, newargs,
+                                            nargs, kwnames);
         newargs[0] = tmp;
     }
     else {
         Py_ssize_t nkwargs = (kwnames == NULL) ? 0 : PyTuple_GET_SIZE(kwnames);
         Py_ssize_t totalargs = nargs + nkwargs;
         if (totalargs == 0) {
-            return _PyObject_Vectorcall(func, &self, 1, NULL);
+            return _PyObject_VectorcallTstate(tstate, func, &self, 1, NULL);
         }
 
         PyObject *newargs_stack[_PY_FASTCALL_SMALL_STACK];
@@ -66,7 +70,7 @@ method_vectorcall(PyObject *method, PyObject *const *args,
         else {
             newargs = PyMem_Malloc((totalargs+1) * sizeof(PyObject *));
             if (newargs == NULL) {
-                PyErr_NoMemory();
+                _PyErr_NoMemory(tstate);
                 return NULL;
             }
         }
@@ -77,7 +81,8 @@ method_vectorcall(PyObject *method, PyObject *const *args,
          * undefined behaviour. */
         assert(args != NULL);
         memcpy(newargs + 1, args, totalargs * sizeof(PyObject *));
-        result = _PyObject_Vectorcall(func, newargs, nargs+1, kwnames);
+        result = _PyObject_VectorcallTstate(tstate, func,
+                                            newargs, nargs+1, kwnames);
         if (newargs != newargs_stack) {
             PyMem_Free(newargs);
         }
diff --git a/Objects/typeobject.c b/Objects/typeobject.c
index 0e1cb7b..50a3c15 100644
--- a/Objects/typeobject.c
+++ b/Objects/typeobject.c
@@ -1445,7 +1445,7 @@ lookup_method(PyObject *self, _Py_Identifier *attrid, int *unbound)
 
 
 static inline PyObject*
-vectorcall_unbound(int unbound, PyObject *func,
+vectorcall_unbound(PyThreadState *tstate, int unbound, PyObject *func,
                    PyObject *const *args, Py_ssize_t nargs)
 {
     size_t nargsf = nargs;
@@ -1455,7 +1455,7 @@ vectorcall_unbound(int unbound, PyObject *func,
         args++;
         nargsf = nargsf - 1 + PY_VECTORCALL_ARGUMENTS_OFFSET;
     }
-    return _PyObject_Vectorcall(func, args, nargsf, NULL);
+    return _PyObject_VectorcallTstate(tstate, func, args, nargsf, NULL);
 }
 
 static PyObject*
@@ -1479,13 +1479,15 @@ vectorcall_method(_Py_Identifier *name,
                   PyObject *const *args, Py_ssize_t nargs)
 {
     assert(nargs >= 1);
+
+    PyThreadState *tstate = _PyThreadState_GET();
     int unbound;
     PyObject *self = args[0];
     PyObject *func = lookup_method(self, name, &unbound);
     if (func == NULL) {
         return NULL;
     }
-    PyObject *retval = vectorcall_unbound(unbound, func, args, nargs);
+    PyObject *retval = vectorcall_unbound(tstate, unbound, func, args, nargs);
     Py_DECREF(func);
     return retval;
 }
@@ -1493,10 +1495,11 @@ vectorcall_method(_Py_Identifier *name,
 /* Clone of vectorcall_method() that returns NotImplemented
  * when the lookup fails. */
 static PyObject *
-vectorcall_maybe(_Py_Identifier *name,
+vectorcall_maybe(PyThreadState *tstate, _Py_Identifier *name,
                  PyObject *const *args, Py_ssize_t nargs)
 {
     assert(nargs >= 1);
+
     int unbound;
     PyObject *self = args[0];
     PyObject *func = lookup_maybe_method(self, name, &unbound);
@@ -1505,7 +1508,7 @@ vectorcall_maybe(_Py_Identifier *name,
             Py_RETURN_NOTIMPLEMENTED;
         return NULL;
     }
-    PyObject *retval = vectorcall_unbound(unbound, func, args, nargs);
+    PyObject *retval = vectorcall_unbound(tstate, unbound, func, args, nargs);
     Py_DECREF(func);
     return retval;
 }
@@ -6177,6 +6180,7 @@ static PyObject * \
 FUNCNAME(PyObject *self, PyObject *other) \
 { \
     PyObject* stack[2]; \
+    PyThreadState *tstate = _PyThreadState_GET(); \
     _Py_static_string(op_id, OPSTR); \
     _Py_static_string(rop_id, ROPSTR); \
     int do_other = Py_TYPE(self) != Py_TYPE(other) && \
@@ -6193,7 +6197,7 @@ FUNCNAME(PyObject *self, PyObject *other) \
             if (ok) { \
                 stack[0] = other; \
                 stack[1] = self; \
-                r = vectorcall_maybe(&rop_id, stack, 2); \
+                r = vectorcall_maybe(tstate, &rop_id, stack, 2); \
                 if (r != Py_NotImplemented) \
                     return r; \
                 Py_DECREF(r); \
@@ -6202,7 +6206,7 @@ FUNCNAME(PyObject *self, PyObject *other) \
         } \
         stack[0] = self; \
         stack[1] = other; \
-        r = vectorcall_maybe(&op_id, stack, 2); \
+        r = vectorcall_maybe(tstate, &op_id, stack, 2); \
         if (r != Py_NotImplemented || \
             Py_TYPE(other) == Py_TYPE(self)) \
             return r; \
@@ -6211,7 +6215,7 @@ FUNCNAME(PyObject *self, PyObject *other) \
     if (do_other) { \
         stack[0] = other; \
         stack[1] = self; \
-        return vectorcall_maybe(&rop_id, stack, 2); \
+        return vectorcall_maybe(tstate, &rop_id, stack, 2); \
     } \
     Py_RETURN_NOTIMPLEMENTED; \
 }
@@ -6293,6 +6297,7 @@ slot_sq_ass_item(PyObject *self, Py_ssize_t index, PyObject *value)
 static int
 slot_sq_contains(PyObject *self, PyObject *value)
 {
+    PyThreadState *tstate = _PyThreadState_GET();
     PyObject *func, *res;
     int result = -1, unbound;
     _Py_IDENTIFIER(__contains__);
@@ -6307,7 +6312,7 @@ slot_sq_contains(PyObject *self, PyObject *value)
     }
     if (func != NULL) {
         PyObject *args[2] = {self, value};
-        res = vectorcall_unbound(unbound, func, args, 2);
+        res = vectorcall_unbound(tstate, unbound, func, args, 2);
         Py_DECREF(func);
         if (res != NULL) {
             result = PyObject_IsTrue(res);
@@ -6682,6 +6687,7 @@ static _Py_Identifier name_op[] = {
 static PyObject *
 slot_tp_richcompare(PyObject *self, PyObject *other, int op)
 {
+    PyThreadState *tstate = _PyThreadState_GET();
     int unbound;
     PyObject *func, *res;
 
@@ -6692,7 +6698,7 @@ slot_tp_richcompare(PyObject *self, PyObject *other, int op)
     }
 
     PyObject *stack[2] = {self, other};
-    res = vectorcall_unbound(unbound, func, stack, 2);
+    res = vectorcall_unbound(tstate, unbound, func, stack, 2);
     Py_DECREF(func);
     return res;
 }
diff --git a/Python/context.c b/Python/context.c
index f48c376..26f2299 100644
--- a/Python/context.c
+++ b/Python/context.c
@@ -3,6 +3,7 @@
 #include "pycore_context.h"
 #include "pycore_hamt.h"
 #include "pycore_object.h"
+#include "pycore_pyerrors.h"
 #include "pycore_pystate.h"
 #include "structmember.h"
 
@@ -101,21 +102,18 @@ PyContext_CopyCurrent(void)
 }
 
 
-int
-PyContext_Enter(PyObject *octx)
+static int
+_PyContext_Enter(PyThreadState *ts, PyObject *octx)
 {
     ENSURE_Context(octx, -1)
     PyContext *ctx = (PyContext *)octx;
 
     if (ctx->ctx_entered) {
-        PyErr_Format(PyExc_RuntimeError,
-                     "cannot enter context: %R is already entered", ctx);
+        _PyErr_Format(ts, PyExc_RuntimeError,
+                      "cannot enter context: %R is already entered", ctx);
         return -1;
     }
 
-    PyThreadState *ts = _PyThreadState_GET();
-    assert(ts != NULL);
-
     ctx->ctx_prev = (PyContext *)ts->context;  /* borrow */
     ctx->ctx_entered = 1;
 
@@ -128,7 +126,16 @@ PyContext_Enter(PyObject *octx)
 
 
 int
-PyContext_Exit(PyObject *octx)
+PyContext_Enter(PyObject *octx)
+{
+    PyThreadState *ts = _PyThreadState_GET();
+    assert(ts != NULL);
+    return _PyContext_Enter(ts, octx);
+}
+
+
+static int
+_PyContext_Exit(PyThreadState *ts, PyObject *octx)
 {
     ENSURE_Context(octx, -1)
     PyContext *ctx = (PyContext *)octx;
@@ -139,9 +146,6 @@ PyContext_Exit(PyObject *octx)
         return -1;
     }
 
-    PyThreadState *ts = _PyThreadState_GET();
-    assert(ts != NULL);
-
     if (ts->context != (PyObject *)ctx) {
         /* Can only happen if someone misuses the C API */
         PyErr_SetString(PyExc_RuntimeError,
@@ -159,6 +163,14 @@ PyContext_Exit(PyObject *octx)
     return 0;
 }
 
+int
+PyContext_Exit(PyObject *octx)
+{
+    PyThreadState *ts = _PyThreadState_GET();
+    assert(ts != NULL);
+    return _PyContext_Exit(ts, octx);
+}
+
 
 PyObject *
 PyContextVar_New(const char *name, PyObject *def)
@@ -621,20 +633,22 @@ static PyObject *
 context_run(PyContext *self, PyObject *const *args,
             Py_ssize_t nargs, PyObject *kwnames)
 {
+    PyThreadState *ts = _PyThreadState_GET();
+
     if (nargs < 1) {
-        PyErr_SetString(PyExc_TypeError,
-                        "run() missing 1 required positional argument");
+        _PyErr_SetString(ts, PyExc_TypeError,
+                         "run() missing 1 required positional argument");
         return NULL;
     }
 
-    if (PyContext_Enter((PyObject *)self)) {
+    if (_PyContext_Enter(ts, (PyObject *)self)) {
         return NULL;
     }
 
-    PyObject *call_result = _PyObject_Vectorcall(
-        args[0], args + 1, nargs - 1, kwnames);
+    PyObject *call_result = _PyObject_VectorcallTstate(
+        ts, args[0], args + 1, nargs - 1, kwnames);
 
-    if (PyContext_Exit((PyObject *)self)) {
+    if (_PyContext_Exit(ts, (PyObject *)self)) {
         return NULL;
     }