add PyModule_Dict as a hook

Summary: Adds `PyModule_Dict` as a hook. Discussed with Carl; looked into the option of removing this hook altogether, but there are a few use cases in `import.c` that still rely on branching at this function: https://fburl.com/code/c6eamlmq (*attempted this here: D51317520*). There might be a better path to upstreamability, but wanted to throw this out for feedback. Reviewed By: carljm Differential Revision: D51320293 fbshipit-source-id: 4fe651e02a5c77967867e977d9d97a771687c686
facebookincubator · Nov 20, 2023 · b5e078b · b5e078b
1 parent e68c565
commit b5e078b
Show file tree

Hide file tree

Showing 7 changed files with 25 additions and 11 deletions.
diff --git a/CinderX/Shadowcode/shadowcode.c b/CinderX/Shadowcode/shadowcode.c
@@ -1567,7 +1567,7 @@ _PyShadow_GetAttrModule(_PyShadow_EvalState *state,
         return 0;
     }
 
-    PyObject *dict = PyModule_Dict(owner);
+    PyObject *dict = Ci_PyModule_Dict(owner);
 
     if (dict) {
 
@@ -1883,7 +1883,7 @@ _PyShadow_LoadMethodFromModule(_PyShadow_EvalState *state,
         return meth_found;
     }
 
-    PyObject *dict = PyModule_Dict(obj);
+    PyObject *dict = Ci_PyModule_Dict(obj);
     if (dict) {
         int version;
         PyObject *value;

diff --git a/CinderX/StaticPython/classloader.c b/CinderX/StaticPython/classloader.c
@@ -4278,7 +4278,7 @@ _PyClassLoader_GetIndirectPtr(PyObject *path, PyObject *func, PyObject *containe
     } else if (PyModule_Check(container)) {
         /* modules have no special translation on things we invoke, so
          * we just rely upon the normal JIT dict watchers */
-        PyObject *dict = PyModule_Dict(container);
+        PyObject *dict = Ci_PyModule_Dict(container);
         if (dict != NULL) {
             cache = _PyJIT_GetDictCache(dict, name);
         }

diff --git a/CinderX/_cinderx.cpp b/CinderX/_cinderx.cpp
@@ -42,6 +42,10 @@ static void shadowcode_code_sizeof(struct _PyShadowCode *shadow, Py_ssize_t *res
     res += sizeof(_Py_CODEUNIT) * shadow->len;
 }
 
+static inline int _PyStrictModule_Check(PyObject* obj) {
+  return PyStrictModule_Check(obj);
+}
+
 static int cinder_init() {
   Ci_hook_type_created = _PyJIT_TypeCreated;
   Ci_hook_type_destroyed = _PyJIT_TypeDestroyed;
@@ -62,6 +66,7 @@ static int cinder_init() {
   Ci_hook_PyJIT_GenYieldFromValue = _PyJIT_GenYieldFromValue;
   Ci_hook_PyJIT_GenMaterializeFrame = _PyJIT_GenMaterializeFrame;
   Ci_hook__PyShadow_FreeAll = _PyShadow_FreeAll;
+  Ci_hook_PyStrictModule_Check = _PyStrictModule_Check;
 
   init_already_existing_types();
 

diff --git a/Include/cinder/hooks.h b/Include/cinder/hooks.h
@@ -89,3 +89,6 @@ typedef PyFrameObject *(*Ci_HookType_PyJIT_GenMaterializeFrame)(
     PyGenObject *gen);
 CiAPI_DATA(Ci_HookType_PyJIT_GenMaterializeFrame)
     Ci_hook_PyJIT_GenMaterializeFrame;
+
+typedef int (*Ci_HookType_PyStrictModule_Check)(PyObject *obj);
+CiAPI_DATA(Ci_HookType_PyStrictModule_Check) Ci_hook_PyStrictModule_Check;
diff --git a/Include/internal/pycore_moduleobject.h b/Include/internal/pycore_moduleobject.h
@@ -1,5 +1,8 @@
 #ifndef Py_INTERNAL_MODULEOBJECT_H
 #define Py_INTERNAL_MODULEOBJECT_H
+
+#include "cinder/hooks.h"
+
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -45,11 +48,12 @@ typedef struct {
     PyObject *imported_from;
 } PyStrictModuleObject;
 
-#ifdef ENABLE_CINDERX
-#   define PyModule_Dict(op) (PyStrictModule_Check(op) ? ((PyStrictModuleObject *)op)->globals : ((PyModuleObject *)op)->md_dict)
-#else
-#   define PyModule_Dict(op) (((PyModuleObject *)op)->md_dict)
-#endif
+static inline PyObject* Ci_PyModule_Dict(PyObject *op) {
+    if (Ci_hook_PyStrictModule_Check && Ci_hook_PyStrictModule_Check(op))
+        return ((PyStrictModuleObject *)op)->globals;
+
+    return ((PyModuleObject *)op)->md_dict;
+}
 
 CiAPI_STATIC_INLINE_FUNC(PyObject*) _PyStrictModuleGetDict(PyObject *mod);
 CiAPI_STATIC_INLINE_FUNC(PyObject*) _PyStrictModuleGetDictSetter(PyObject *mod);

diff --git a/Python/cinderhooks.c b/Python/cinderhooks.c
@@ -34,3 +34,5 @@ Ci_HookType_PyJIT_GenDealloc Ci_hook_PyJIT_GenDealloc = NULL;
 Ci_HookType_PyJIT_GenSend Ci_hook_PyJIT_GenSend = NULL;
 Ci_HookType_PyJIT_GenYieldFromValue Ci_hook_PyJIT_GenYieldFromValue = NULL;
 Ci_HookType_PyJIT_GenMaterializeFrame  Ci_hook_PyJIT_GenMaterializeFrame = NULL;
+
+Ci_HookType_PyStrictModule_Check Ci_hook_PyStrictModule_Check = NULL;
diff --git a/Python/import.c b/Python/import.c
@@ -8,7 +8,7 @@
 #include "pycore_interp.h"        // _PyInterpreterState_ClearModules()
 #include "pycore_lazyimport.h"    // PyLazyImport_CheckExact()
 #include "pycore_long.h"          // _PyLong_GetZero()
-#include "pycore_moduleobject.h"  // PyModule_Dict()
+#include "pycore_moduleobject.h"  // Ci_PyModule_Dict()
 #include "pycore_pyerrors.h"
 #include "pycore_pyhash.h"
 #include "pycore_pylifecycle.h"
@@ -2975,7 +2975,7 @@ _imp_hydrate_lazy_objects_impl(PyObject *module)
             if (dst_module != Py_None) {
                 Py_XDECREF(dst_dict);
                 if (PyModule_Check(dst_module)) {
-                    dst_dict = PyModule_Dict(dst_module);
+                    dst_dict = Ci_PyModule_Dict(dst_module);
                     Py_XINCREF(dst_dict);
                 } else {
                     dst_dict = _PyObject_GetAttrId(dst_module, &PyId___dict__);
@@ -3000,7 +3000,7 @@ _imp_hydrate_lazy_objects_impl(PyObject *module)
                             if (d->lz_attr) {
                                 Py_XDECREF(src_dict);
                                 if (PyModule_Check(src_module)) {
-                                    src_dict = PyModule_Dict(src_module);
+                                    src_dict = Ci_PyModule_Dict(src_module);
                                     Py_XINCREF(src_dict);
                                 } else {
                                     src_dict = _PyObject_GetAttrId(src_module, &PyId___dict__);