gh-101476: Add _PyType_GetModuleState (GH-101477)

For fast module state access from heap type methods.

diff --git a/Include/internal/pycore_typeobject.h b/Include/internal/pycore_typeobject.h
index 4d705740a9a62b932b3490f519a1d32cfe520c64..cc5ce2875101ea838dab5e445707f3a437618fc1 100644 (file)
--- a/Include/internal/pycore_typeobject.h
+++ b/Include/internal/pycore_typeobject.h
@@ -4,6 +4,8 @@
 extern "C" {
 #endif
 
+#include "pycore_moduleobject.h"
+
 #ifndef Py_BUILD_CORE
 #  error "this header requires Py_BUILD_CORE define"
 #endif
@@ -62,6 +64,20 @@ _PyStaticType_GET_WEAKREFS_LISTPTR(static_builtin_state *state)
     return &state->tp_weaklist;
 }
 
+/* Like PyType_GetModuleState, but skips verification
+ * that type is a heap type with an associated module */
+static inline void *
+_PyType_GetModuleState(PyTypeObject *type)
+{
+    assert(PyType_Check(type));
+    assert(type->tp_flags & Py_TPFLAGS_HEAPTYPE);
+    PyHeapTypeObject *et = (PyHeapTypeObject *)type;
+    assert(et->ht_module);
+    PyModuleObject *mod = (PyModuleObject *)(et->ht_module);
+    assert(mod != NULL);
+    return mod->md_state;
+}
+
 struct types_state {
     struct type_cache type_cache;
     size_t num_builtins_initialized;

diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c
index 6986695e47b1aec715b762977d4e1acde1e46948..c986e02867ca82d2c8ebb935899fc71f0814e95c 100644 (file)
--- a/Modules/itertoolsmodule.c
+++ b/Modules/itertoolsmodule.c
@@ -3,6 +3,7 @@
 #include "pycore_call.h"          // _PyObject_CallNoArgs()
 #include "pycore_long.h"          // _PyLong_GetZero()
 #include "pycore_moduleobject.h"  // _PyModule_GetState()
+#include "pycore_typeobject.h"    // _PyType_GetModuleState()
 #include "pycore_object.h"        // _PyObject_GC_TRACK()
 #include "pycore_tuple.h"         // _PyTuple_ITEMS()
 #include "structmember.h"         // PyMemberDef
@@ -48,7 +49,7 @@ get_module_state(PyObject *mod)
 static inline itertools_state *
 get_module_state_by_cls(PyTypeObject *cls)
 {
-    void *state = PyType_GetModuleState(cls);
+    void *state = _PyType_GetModuleState(cls);
     assert(state != NULL);
     return (itertools_state *)state;
 }