gh-101476: Add _PyType_GetModuleState (GH-101477)

For fast module state access from heap type methods.

Author: erlend-aasland

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;

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;
 }