1 files changed, 42 insertions, 15 deletions

diff --git a/Python/pystate.c b/Python/pystate.c
index 2ac2fd5274..83f15fd671 100644
--- a/Python/pystate.c
+++ b/Python/pystate.c
@@ -3,6 +3,12 @@
 
 #include "Python.h"
 
+#ifndef Py_BUILD_CORE
+/* ensure that PyThreadState_GET() is a macro, not an alias to
+ * PyThreadState_Get() */
+#  error "pystate.c must be compiled with Py_BUILD_CORE defined"
+#endif
+
 /* --------------------------------------------------------------------------
 CAUTION
 
@@ -212,6 +218,9 @@ new_threadstate(PyInterpreterState *interp, int init)
         tstate->on_delete = NULL;
         tstate->on_delete_data = NULL;
 
+        tstate->coroutine_wrapper = NULL;
+        tstate->in_coroutine_wrapper = 0;
+
         if (init)
             _PyThreadState_Init(tstate);
 
@@ -253,6 +262,9 @@ PyState_FindModule(struct PyModuleDef* module)
     Py_ssize_t index = module->m_base.m_index;
     PyInterpreterState *state = PyThreadState_GET()->interp;
     PyObject *res;
+    if (module->m_slots) {
+        return NULL;
+    }
     if (index == 0)
         return NULL;
     if (state->modules_by_index == NULL)
@@ -266,7 +278,13 @@ PyState_FindModule(struct PyModuleDef* module)
 int
 _PyState_AddModule(PyObject* module, struct PyModuleDef* def)
 {
-    PyInterpreterState *state = PyThreadState_GET()->interp;
+    PyInterpreterState *state;
+    if (def->m_slots) {
+        PyErr_SetString(PyExc_SystemError,
+                        "PyState_AddModule called on module with slots");
+        return -1;
+    }
+    state = PyThreadState_GET()->interp;
     if (!def)
         return -1;
     if (!state->modules_by_index) {
@@ -306,8 +324,14 @@ PyState_AddModule(PyObject* module, struct PyModuleDef* def)
 int
 PyState_RemoveModule(struct PyModuleDef* def)
 {
+    PyInterpreterState *state;
     Py_ssize_t index = def->m_base.m_index;
-    PyInterpreterState *state = PyThreadState_GET()->interp;
+    if (def->m_slots) {
+        PyErr_SetString(PyExc_SystemError,
+                        "PyState_RemoveModule called on module with slots");
+        return -1;
+    }
+    state = PyThreadState_GET()->interp;
     if (index == 0) {
         Py_FatalError("PyState_RemoveModule: Module index invalid.");
         return -1;
@@ -372,6 +396,8 @@ PyThreadState_Clear(PyThreadState *tstate)
     tstate->c_tracefunc = NULL;
     Py_CLEAR(tstate->c_profileobj);
     Py_CLEAR(tstate->c_traceobj);
+
+    Py_CLEAR(tstate->coroutine_wrapper);
 }
 
 
@@ -403,7 +429,7 @@ tstate_delete_common(PyThreadState *tstate)
 void
 PyThreadState_Delete(PyThreadState *tstate)
 {
-    if (tstate == _Py_atomic_load_relaxed(&_PyThreadState_Current))
+    if (tstate == PyThreadState_GET())
         Py_FatalError("PyThreadState_Delete: tstate is still current");
 #ifdef WITH_THREAD
     if (autoInterpreterState && PyThread_get_key_value(autoTLSkey) == tstate)
@@ -417,8 +443,7 @@ PyThreadState_Delete(PyThreadState *tstate)
 void
 PyThreadState_DeleteCurrent()
 {
-    PyThreadState *tstate = (PyThreadState*)_Py_atomic_load_relaxed(
-        &_PyThreadState_Current);
+    PyThreadState *tstate = PyThreadState_GET();
     if (tstate == NULL)
         Py_FatalError(
             "PyThreadState_DeleteCurrent: no current tstate");
@@ -469,10 +494,16 @@ _PyThreadState_DeleteExcept(PyThreadState *tstate)
 
 
 PyThreadState *
+_PyThreadState_UncheckedGet(void)
+{
+    return PyThreadState_GET();
+}
+
+
+PyThreadState *
 PyThreadState_Get(void)
 {
-    PyThreadState *tstate = (PyThreadState*)_Py_atomic_load_relaxed(
-        &_PyThreadState_Current);
+    PyThreadState *tstate = PyThreadState_GET();
     if (tstate == NULL)
         Py_FatalError("PyThreadState_Get: no current thread");
 
@@ -483,8 +514,7 @@ PyThreadState_Get(void)
 PyThreadState *
 PyThreadState_Swap(PyThreadState *newts)
 {
-    PyThreadState *oldts = (PyThreadState*)_Py_atomic_load_relaxed(
-        &_PyThreadState_Current);
+    PyThreadState *oldts = PyThreadState_GET();
 
     _Py_atomic_store_relaxed(&_PyThreadState_Current, newts);
     /* It should not be possible for more than one thread state
@@ -515,8 +545,7 @@ PyThreadState_Swap(PyThreadState *newts)
 PyObject *
 PyThreadState_GetDict(void)
 {
-    PyThreadState *tstate = (PyThreadState*)_Py_atomic_load_relaxed(
-        &_PyThreadState_Current);
+    PyThreadState *tstate = PyThreadState_GET();
     if (tstate == NULL)
         return NULL;
 
@@ -662,7 +691,7 @@ PyThreadState_IsCurrent(PyThreadState *tstate)
 {
     /* Must be the tstate for this thread */
     assert(PyGILState_GetThisThreadState()==tstate);
-    return tstate == _Py_atomic_load_relaxed(&_PyThreadState_Current);
+    return tstate == PyThreadState_GET();
 }
 
 /* Internal initialization/finalization functions called by
@@ -754,9 +783,7 @@ PyGILState_GetThisThreadState(void)
 int
 PyGILState_Check(void)
 {
-    /* can't use PyThreadState_Get() since it will assert that it has the GIL */
-    PyThreadState *tstate = (PyThreadState*)_Py_atomic_load_relaxed(
-        &_PyThreadState_Current);
+    PyThreadState *tstate = PyThreadState_GET();
     return tstate && (tstate == PyGILState_GetThisThreadState());
 }