4 files changed, 26 insertions, 46 deletions

diff --git a/Include/internal/pycore_typeobject.h b/Include/internal/pycore_typeobject.h
index ca5a1e2..118bc98 100644
--- a/Include/internal/pycore_typeobject.h
+++ b/Include/internal/pycore_typeobject.h
@@ -209,7 +209,6 @@ extern PyObject * _PyType_GetBases(PyTypeObject *type);
 extern PyObject * _PyType_GetMRO(PyTypeObject *type);
 extern PyObject* _PyType_GetSubclasses(PyTypeObject *);
 extern int _PyType_HasSubclasses(PyTypeObject *);
-PyAPI_FUNC(PyObject *) _PyType_GetModuleByDef2(PyTypeObject *, PyTypeObject *, PyModuleDef *);
 
 // Export for _testinternalcapi extension.
 PyAPI_FUNC(PyObject *) _PyType_GetSlotWrapperNames(void);
diff --git a/Modules/_collectionsmodule.c b/Modules/_collectionsmodule.c
index fbfed59..aef0424 100644
--- a/Modules/_collectionsmodule.c
+++ b/Modules/_collectionsmodule.c
@@ -2179,6 +2179,8 @@ typedef struct {
     PyObject *default_factory;
 } defdictobject;
 
+static PyType_Spec defdict_spec;
+
 PyDoc_STRVAR(defdict_missing_doc,
 "__missing__(key) # Called by __getitem__ for missing key; pseudo-code:\n\
   if self.default_factory is None: raise KeyError((key,))\n\
@@ -2358,23 +2360,16 @@ defdict_or(PyObject* left, PyObject* right)
 {
     PyObject *self, *other;
 
-    // Find module state
-    PyTypeObject *tp = Py_TYPE(left);
-    PyObject *mod = PyType_GetModuleByDef(tp, &_collectionsmodule);
-    if (mod == NULL) {
-        PyErr_Clear();
-        tp = Py_TYPE(right);
-        mod = PyType_GetModuleByDef(tp, &_collectionsmodule);
+    int ret = PyType_GetBaseByToken(Py_TYPE(left), &defdict_spec, NULL);
+    if (ret < 0) {
+        return NULL;
     }
-    assert(mod != NULL);
-    collections_state *state = get_module_state(mod);
-
-    if (PyObject_TypeCheck(left, state->defdict_type)) {
+    if (ret) {
         self = left;
         other = right;
     }
     else {
-        assert(PyObject_TypeCheck(right, state->defdict_type));
+        assert(PyType_GetBaseByToken(Py_TYPE(right), &defdict_spec, NULL) == 1);
         self = right;
         other = left;
     }
@@ -2454,6 +2449,7 @@ passed to the dict constructor, including keyword arguments.\n\
 #define DEFERRED_ADDRESS(ADDR) 0
 
 static PyType_Slot defdict_slots[] = {
+    {Py_tp_token, Py_TP_USE_SPEC},
     {Py_tp_dealloc, defdict_dealloc},
     {Py_tp_repr, defdict_repr},
     {Py_nb_or, defdict_or},
diff --git a/Modules/_decimal/_decimal.c b/Modules/_decimal/_decimal.c
index e99a96a..68d1da9 100644
--- a/Modules/_decimal/_decimal.c
+++ b/Modules/_decimal/_decimal.c
@@ -122,6 +122,7 @@ get_module_state(PyObject *mod)
 }
 
 static struct PyModuleDef _decimal_module;
+static PyType_Spec dec_spec;
 
 static inline decimal_state *
 get_module_state_by_def(PyTypeObject *tp)
@@ -134,10 +135,16 @@ get_module_state_by_def(PyTypeObject *tp)
 static inline decimal_state *
 find_state_left_or_right(PyObject *left, PyObject *right)
 {
-    PyObject *mod = _PyType_GetModuleByDef2(Py_TYPE(left), Py_TYPE(right),
-                                            &_decimal_module);
-    assert(mod != NULL);
-    return get_module_state(mod);
+    PyTypeObject *base;
+    if (PyType_GetBaseByToken(Py_TYPE(left), &dec_spec, &base) != 1) {
+        assert(!PyErr_Occurred());
+        PyType_GetBaseByToken(Py_TYPE(right), &dec_spec, &base);
+    }
+    assert(base != NULL);
+    void *state = _PyType_GetModuleState(base);
+    assert(state != NULL);
+    Py_DECREF(base);
+    return (decimal_state *)state;
 }
 
 
@@ -745,7 +752,7 @@ signaldict_richcompare(PyObject *v, PyObject *w, int op)
 {
     PyObject *res = Py_NotImplemented;
 
-    decimal_state *state = find_state_left_or_right(v, w);
+    decimal_state *state = get_module_state_by_def(Py_TYPE(v));
     assert(PyDecSignalDict_Check(state, v));
 
     if ((SdFlagAddr(v) == NULL) || (SdFlagAddr(w) == NULL)) {
@@ -5041,6 +5048,7 @@ static PyMethodDef dec_methods [] =
 };
 
 static PyType_Slot dec_slots[] = {
+    {Py_tp_token, Py_TP_USE_SPEC},
     {Py_tp_dealloc, dec_dealloc},
     {Py_tp_getattro, PyObject_GenericGetAttr},
     {Py_tp_traverse, dec_traverse},
diff --git a/Objects/typeobject.c b/Objects/typeobject.c
index 68e481f..3368c1e 100644
--- a/Objects/typeobject.c
+++ b/Objects/typeobject.c
@@ -5207,8 +5207,8 @@ PyType_GetModuleState(PyTypeObject *type)
 /* Get the module of the first superclass where the module has the
  * given PyModuleDef.
  */
-static inline PyObject *
-get_module_by_def(PyTypeObject *type, PyModuleDef *def)
+PyObject *
+PyType_GetModuleByDef(PyTypeObject *type, PyModuleDef *def)
 {
     assert(PyType_Check(type));
 
@@ -5241,7 +5241,7 @@ get_module_by_def(PyTypeObject *type, PyModuleDef *def)
     Py_ssize_t n = PyTuple_GET_SIZE(mro);
     for (Py_ssize_t i = 1; i < n; i++) {
         PyObject *super = PyTuple_GET_ITEM(mro, i);
-        if(!_PyType_HasFeature((PyTypeObject *)super, Py_TPFLAGS_HEAPTYPE)) {
+        if (!_PyType_HasFeature((PyTypeObject *)super, Py_TPFLAGS_HEAPTYPE)) {
             // Static types in the MRO need to be skipped
             continue;
         }
@@ -5254,37 +5254,14 @@ get_module_by_def(PyTypeObject *type, PyModuleDef *def)
         }
     }
     END_TYPE_LOCK();
-    return res;
-}
 
-PyObject *
-PyType_GetModuleByDef(PyTypeObject *type, PyModuleDef *def)
-{
-    PyObject *module = get_module_by_def(type, def);
-    if (module == NULL) {
+    if (res == NULL) {
         PyErr_Format(
             PyExc_TypeError,
             "PyType_GetModuleByDef: No superclass of '%s' has the given module",
             type->tp_name);
     }
-    return module;
-}
-
-PyObject *
-_PyType_GetModuleByDef2(PyTypeObject *left, PyTypeObject *right,
-                        PyModuleDef *def)
-{
-    PyObject *module = get_module_by_def(left, def);
-    if (module == NULL) {
-        module = get_module_by_def(right, def);
-        if (module == NULL) {
-            PyErr_Format(
-                PyExc_TypeError,
-                "PyType_GetModuleByDef: No superclass of '%s' nor '%s' has "
-                "the given module", left->tp_name, right->tp_name);
-        }
-    }
-    return module;
+    return res;
 }