From ccd98a3146d66343499d04a44e038223a1a09e80 Mon Sep 17 00:00:00 2001
From: "Erlend E. Aasland" <erlend.aasland@protonmail.com>
Date: Thu, 23 Feb 2023 22:42:15 +0100
Subject: [PATCH] gh-101476: Add _PyType_GetModuleState (GH-101477)

For fast module state access from heap type methods.
---
 Include/internal/pycore_typeobject.h | 16 ++++++++++++++++
 Modules/itertoolsmodule.c            |  3 ++-
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/Include/internal/pycore_typeobject.h b/Include/internal/pycore_typeobject.h
index 4d705740a9a..cc5ce287510 100644
--- 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 6986695e47b..c986e02867c 100644
--- 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;
 }