RFC: Basic module state support

hpy/debug/src/debug_ctx_cpython.c

@@ -24,6 +24,7 @@
 #include <Python.h>
 #include "debug_internal.h"
 #include "hpy/runtime/ctx_type.h" // for call_traverseproc_from_trampoline
+#include "hpy/runtime/ctx_module.h" // for call_mod_traverseproc_from_trampoline
 #include "handles.h" // for _py2h and _h2py
 #if defined(_MSC_VER)
 # include <malloc.h>   /* for alloca() */
@@ -186,6 +187,13 @@ void debug_ctx_CallRealFunctionFromTrampoline(HPyContext *dctx,
                                                       a->visit, a->arg);
         return;
     }
+    case HPyFunc_MODTRAVERSEPROC: {
+        HPyFunc_traverseproc f = (HPyFunc_traverseproc)func;
+        _HPyFunc_args_TRAVERSEPROC *a = (_HPyFunc_args_TRAVERSEPROC*)args;
+        a->result = call_mod_traverseproc_from_trampoline(f, a->self,
+                                                          a->visit, a->arg);
+        return;
+    }
 #include "autogen_debug_ctx_call.i"
     default:
         Py_FatalError("Unsupported HPyFunc_Signature in debug_ctx_cpython.c");

hpy/devel/include/hpy.h

@@ -168,6 +168,7 @@ typedef struct _HPyContext_s HPyContext;
 //  CPython-ABI
 #   include "hpy/runtime/ctx_funcs.h"
 #   include "hpy/runtime/ctx_type.h"
+#   include "hpy/runtime/ctx_module.h"
 #   include "hpy/cpython/misc.h"
 #   include "hpy/cpython/autogen_api_impl.h"
 #endif

hpy/devel/include/hpy/cpython/hpyfunc_trampolines.h

@@ -68,6 +68,10 @@ typedef int (*_HPyCFunction_INITPROC)(HPyContext*, HPy, HPy *, HPy_ssize_t, HPy)
 #define _HPyFunc_TRAMPOLINE_HPyFunc_DESTROYFUNC(SYM, IMPL)              \
     static void SYM(void) { abort(); }
 
+/* special case: the same for HPy_m_destroy for module state */
+#define _HPyFunc_TRAMPOLINE_HPyFunc_MODDESTROYFUNC(SYM, IMPL)              \
+    static void SYM(void) { abort(); }
+
 /* this needs to be written manually because HPy has a different type for
    "op": HPy_RichCmpOp instead of int */
 typedef HPy (*_HPyCFunction_RICHCMPFUNC)(HPyContext *, HPy, HPy, int);
@@ -107,4 +111,11 @@ typedef int (*_HPyCFunction_RELEASEBUFFERPROC)(HPyContext *, HPy, HPy_buffer *);
                                                  visit, arg);           \
     }
 
+#define _HPyFunc_TRAMPOLINE_HPyFunc_MODTRAVERSEPROC(SYM, IMPL)          \
+    static int SYM(cpy_PyObject *self, cpy_visitproc visit, void *arg)  \
+    {                                                                   \
+        return call_mod_traverseproc_from_trampoline((HPyFunc_traverseproc)IMPL, self,        \
+                                                 visit, arg);           \
+    }
+
 #endif // HPY_CPYTHON_HPYFUNC_TRAMPOLINES_H

hpy/devel/include/hpy/cpython/misc.h

@@ -284,6 +284,11 @@ HPyAPI_FUNC HPy HPyModule_Create(HPyContext *ctx, HPyModuleDef *mdef)
     return ctx_Module_Create(ctx, mdef);
 }
 
+HPyAPI_FUNC void* HPyModule_GetState(HPyContext *ctx, HPy mod)
+{
+    return ctx_Module_GetState(ctx, mod);
+}
+
 HPyAPI_FUNC HPy HPyType_FromSpec(HPyContext *ctx, HPyType_Spec *spec, HPyType_SpecParam *params)
 {
     return ctx_Type_FromSpec(ctx, spec, params);

hpy/devel/include/hpy/hpyfunc.h

@@ -39,6 +39,8 @@ typedef enum {
     HPyFunc_OBJOBJPROC,
     HPyFunc_TRAVERSEPROC,
     HPyFunc_DESTRUCTOR,
+    HPyFunc_MODTRAVERSEPROC,
+    HPyFunc_MODDESTROYFUNC,
 
 } HPyFunc_Signature;
 
@@ -109,6 +111,10 @@ typedef int (*HPyFunc_visitproc)(HPyField *, void *);
 
 #include "autogen_hpyfunc_declare.h"
 
+// Few special cases not handled by the autogen:
+#define _HPyFunc_DECLARE_HPyFunc_MODTRAVERSEPROC(SYM) static int SYM(void *object, HPyFunc_visitproc visit, void *arg)
+#define _HPyFunc_DECLARE_HPyFunc_MODDESTROYFUNC(SYM) static void SYM(void *object)
+
 #ifdef HPY_UNIVERSAL_ABI
 #  include "universal/hpyfunc_trampolines.h"
 #  include "universal/autogen_hpyfunc_trampolines.h"

hpy/devel/include/hpy/runtime/ctx_funcs.h

@@ -26,6 +26,8 @@ _HPy_HIDDEN void ctx_ListBuilder_Cancel(HPyContext *ctx, HPyListBuilder builder)
 
 // ctx_module.c
 _HPy_HIDDEN HPy ctx_Module_Create(HPyContext *ctx, HPyModuleDef *hpydef);
+// ctx_bytes.c
+_HPy_HIDDEN void* ctx_Module_GetState(HPyContext *ctx, HPy mod);
 
 // ctx_object.c
 _HPy_HIDDEN void ctx_Dump(HPyContext *ctx, HPy h);

hpy/devel/include/hpy/runtime/ctx_module.h

@@ -0,0 +1,13 @@
+#ifndef HPY_COMMON_RUNTIME_CTX_MODULE_H
+#define HPY_COMMON_RUNTIME_CTX_MODULE_H
+
+#include <Python.h>
+#include "hpy.h"
+#include "hpy/hpytype.h"
+
+_HPy_HIDDEN int call_mod_traverseproc_from_trampoline(HPyFunc_traverseproc tp_traverse,
+                                                      PyObject *self,
+                                                      cpy_visitproc cpy_visit,
+                                                      void *cpy_arg);
+
+#endif /* HPY_COMMON_RUNTIME_CTX_MODULE_H */

hpy/devel/include/hpy/runtime/ctx_type.h

@@ -8,6 +8,11 @@
 _HPy_HIDDEN PyMethodDef *create_method_defs(HPyDef *hpydefs[],
                                             PyMethodDef *legacy_methods);
 
+_HPy_HIDDEN int call_struct_traverseproc_from_trampoline(HPyFunc_traverseproc tp_traverse,
+                                                         void *struct_data,
+                                                         cpy_visitproc cpy_visit,
+                                                         void *cpy_arg);
+
 _HPy_HIDDEN int call_traverseproc_from_trampoline(HPyFunc_traverseproc tp_traverse,
                                                   PyObject *self,
                                                   cpy_visitproc cpy_visit,

hpy/devel/include/hpy/universal/hpyfunc_trampolines.h

@@ -100,6 +100,9 @@ typedef struct {
 #define _HPyFunc_TRAMPOLINE_HPyFunc_DESTROYFUNC(SYM, IMPL)              \
     static void SYM(void) { abort(); }
 
+/* special case: the same for HPy_m_destroy for module state */
+#define _HPyFunc_TRAMPOLINE_HPyFunc_MODDESTROYFUNC(SYM, IMPL)              \
+    static void SYM(void) { abort(); }
 
 /* this needs to be written manually because HPy has a different type for
    "op": HPy_RichCmpOp instead of int */
@@ -160,6 +163,15 @@ typedef struct {
         return a.result; \
     }
 
+#define _HPyFunc_TRAMPOLINE_HPyFunc_MODTRAVERSEPROC(SYM, IMPL) \
+    static int SYM(cpy_PyObject *self, cpy_visitproc visit, void *arg) \
+    { \
+        _HPyFunc_args_TRAVERSEPROC a = { self, visit, arg }; \
+        _HPy_CallRealFunctionFromTrampoline( \
+           _ctx_for_trampolines, HPyFunc_MODTRAVERSEPROC, (HPyCFunction)IMPL, &a); \
+        return a.result; \
+    }
+
 
 
 #endif // HPY_UNIVERSAL_HPYFUNC_TRAMPOLINES_H

hpy/devel/src/runtime/ctx_module.c

@@ -11,6 +11,53 @@ static PyModuleDef empty_moduledef = {
     PyModuleDef_HEAD_INIT
 };
 
+/* This is a hack similar to what we do with types:
+   We need some extra space to store our extra slots on
+   the modules, the actual user data will be appended
+   to this struct which we must align by using the union
+   at the end.
+*/
+typedef struct {
+    HPyFunc_traverseproc m_traverse_impl;
+    HPyFunc_destroyfunc m_destroy_impl;
+    union {
+        unsigned char user_data_payload[1];
+        max_align_t _align;
+    };
+} HPyMod_HEAD_t;
+
+#define HPyMod_HEAD_SIZE (offsetof(HPyMod_HEAD_t, user_data_payload))
+
+static int _decref_visitor(HPyField *pf, void *arg)
+{
+    PyObject *old_object = _hf2py(*pf);
+    *pf = HPyField_NULL;
+    Py_XDECREF(old_object);
+    return 0;
+}
+
+static int hpymod_clear(PyObject *self)
+{
+    HPyMod_HEAD_t *extra = (HPyMod_HEAD_t*) PyModule_GetState(self);
+    extra->m_traverse_impl(&extra->user_data_payload, _decref_visitor, NULL);
+    return 0;
+}
+
+static void hpymod_free(void *self)
+{
+    HPyMod_HEAD_t *extra = (HPyMod_HEAD_t*) PyModule_GetState((PyObject*) self);
+    if (extra->m_traverse_impl)
+        extra->m_traverse_impl(&extra->user_data_payload, _decref_visitor, NULL);
+    if (extra->m_destroy_impl)
+        extra->m_destroy_impl(&extra->user_data_payload);
+}
+
+_HPy_HIDDEN void*
+ctx_Module_GetState(HPyContext *ctx, HPy h_mod) {
+    HPyMod_HEAD_t *internal = (HPyMod_HEAD_t *) PyModule_GetState(_h2py(h_mod));
+    return internal ? &internal->user_data_payload : NULL;
+}
+
 _HPy_HIDDEN HPy
 ctx_Module_Create(HPyContext *ctx, HPyModuleDef *hpydef)
 {
@@ -27,12 +74,63 @@ ctx_Module_Create(HPyContext *ctx, HPyModuleDef *hpydef)
     memcpy(def, &empty_moduledef, sizeof(PyModuleDef));
     def->m_name = hpydef->name;
     def->m_doc = hpydef->doc;
-    def->m_size = hpydef->size;
     def->m_methods = create_method_defs(hpydef->defines, hpydef->legacy_methods);
+
+    HPyMod_HEAD_t extra = { .m_traverse_impl = NULL, .m_destroy_impl = NULL };
+    bool needs_state = hpydef->size != -1;
+    if (hpydef->defines != NULL) {
+        for (int i = 0; hpydef->defines[i] != NULL; i++) {
+            HPyDef *src = hpydef->defines[i];
+            if (src->kind != HPyDef_Kind_Slot)
+                continue;
+            if (src->slot.slot == HPy_m_destroy) {
+                extra.m_destroy_impl = (HPyFunc_destroyfunc) src->slot.impl;
+                // No trampoline, this is HPy specific slot not called by
+                // CPython runtime, but by us in hpymod_free
+                continue;
+            }
+            if (src->slot.slot == HPy_m_traverse) {
+                def->m_traverse = (traverseproc) src->slot.cpy_trampoline;
+                extra.m_traverse_impl = (HPyFunc_traverseproc) src->slot.impl;
+                def->m_clear = &hpymod_clear;
+            } else {
+                const size_t buffer_size = 256;
+                char buffer[buffer_size];
+                snprintf(buffer, buffer_size, "Unsupported slot number %d for module '%s'", src->slot.slot, hpydef->name);
+                HPy_FatalError(ctx, buffer);
+            }
+        }
+    }
+
+    if (needs_state) {
+        // The assumption is such that most of the time if there is a module
+        // state, then there will be some HPyFields to traverse, so we always
+        // allocate our custom state. This way HPyModule_GetState can assume
+        // that state is always HPyMod_Extra_t.
+        def->m_size = HPyMod_HEAD_SIZE + hpydef->size;
+        def->m_free = &hpymod_free;
+    }
+
     if (def->m_methods == NULL) {
         PyMem_Free(def);
         return HPy_NULL;
     }
     PyObject *result = PyModule_Create(def);
+    if (needs_state) {
+        HPyMod_HEAD_t *state = (HPyMod_HEAD_t*) PyModule_GetState(result);
+        state->m_traverse_impl = extra.m_traverse_impl;
+        state->m_destroy_impl = extra.m_destroy_impl;
+    }
     return _py2h(result);
 }
+
+_HPy_HIDDEN int call_mod_traverseproc_from_trampoline(HPyFunc_traverseproc tp_traverse,
+                                                      PyObject *self,
+                                                      cpy_visitproc cpy_visit,
+                                                      void *cpy_arg)
+{
+    return call_struct_traverseproc_from_trampoline(tp_traverse,
+                                                    &((HPyMod_HEAD_t*) PyModule_GetState(self))->user_data_payload,
+                                                    cpy_visit,
+                                                    cpy_arg);
+}

hpy/devel/src/runtime/ctx_type.c

@@ -921,11 +921,22 @@ static int hpy2cpy_visit(HPyField *f, void *v_args)
     return cpy_visit(cpy_obj, cpy_arg);
 }
 
+_HPy_HIDDEN int call_struct_traverseproc_from_trampoline(HPyFunc_traverseproc tp_traverse,
+                                                  void *struct_data,
+                                                  cpy_visitproc cpy_visit,
+                                                  void *cpy_arg)
+{
+    hpy2cpy_visit_args_t args = { cpy_visit, cpy_arg };
+    return tp_traverse(struct_data, hpy2cpy_visit, &args);
+}
+
 _HPy_HIDDEN int call_traverseproc_from_trampoline(HPyFunc_traverseproc tp_traverse,
                                                   PyObject *self,
                                                   cpy_visitproc cpy_visit,
                                                   void *cpy_arg)
 {
-    hpy2cpy_visit_args_t args = { cpy_visit, cpy_arg };
-    return tp_traverse(_pyobj_as_struct(self), hpy2cpy_visit, &args);
+    return call_struct_traverseproc_from_trampoline(tp_traverse,
+                                                    _pyobj_as_struct(self),
+                                                    cpy_visit,
+                                                    cpy_arg);
 }

hpy/tools/autogen/conf.py

@@ -113,4 +113,5 @@
     'HPy_ReenterPythonExecution': 'PyEval_RestoreThread',
     'HPyGlobal_Load': None,
     'HPyGlobal_Store': None,
+    'HPyModule_GetState': None,
 }

hpy/tools/autogen/public_api.h

@@ -114,6 +114,8 @@ HPy h_TupleType;        /* built-in 'tuple' */
 HPy h_ListType;         /* built-in 'list' */
 
 HPy HPyModule_Create(HPyContext *ctx, HPyModuleDef *def);
+void* HPyModule_GetState(HPyContext *ctx, HPy module);
+
 HPy HPy_Dup(HPyContext *ctx, HPy h);
 void HPy_Close(HPyContext *ctx, HPy h);
 
@@ -586,5 +588,7 @@ typedef enum {
 
     /* extra HPy slots */
     HPy_tp_destroy = SLOT(1000, HPyFunc_DESTROYFUNC),
+    HPy_m_traverse = SLOT(1001, HPyFunc_MODTRAVERSEPROC),
+    HPy_m_destroy = SLOT(1002, HPyFunc_MODDESTROYFUNC),
 
 } HPySlot_Slot;