Merge pull request #19404 from mattip/nep49-3

NEP: update NEP with the PyDataMem_Handler struct as implemented in the PR

1 files changed, 62 insertions, 81 deletions

diff --git a/doc/neps/nep-0049.rst b/doc/neps/nep-0049.rst
index 743dd2ad6..4a5edae81 100644
--- a/doc/neps/nep-0049.rst
+++ b/doc/neps/nep-0049.rst
@@ -118,20 +118,31 @@ NumPy C-API functions
 
         typedef struct {
             char name[128];  /* multiple of 64 to keep the struct aligned */
-            PyDataMem_AllocFunc *alloc;
-            PyDataMem_ZeroedAllocFunc *zeroed_alloc;
-            PyDataMem_FreeFunc *free;
-            PyDataMem_ReallocFunc *realloc;
+            PyDataMemAllocator allocator;
         } PyDataMem_Handler;
 
-    where the function's signatures are
+    where the allocator structure is
 
     .. code-block:: c
 
-        typedef void *(PyDataMem_AllocFunc)(size_t size);
-        typedef void *(PyDataMem_ZeroedAllocFunc)(size_t nelems, size_t elsize);
-        typedef void (PyDataMem_FreeFunc)(void *ptr, size_t size);
-        typedef void *(PyDataMem_ReallocFunc)(void *ptr, size_t size);
+        /* The declaration of free differs from PyMemAllocatorEx */
+        typedef struct {
+            void *ctx;
+            void* (*malloc) (void *ctx, size_t size);
+            void* (*calloc) (void *ctx, size_t nelem, size_t elsize);
+            void* (*realloc) (void *ctx, void *ptr, size_t new_size);
+            void (*free) (void *ctx, void *ptr, size_t size);
+        } PyDataMemAllocator;
+
+    The use of a ``size`` parameter in ``free`` differentiates this struct from
+    the :c:type:`PyMemAllocatorEx` struct in Python. This call signature is
+    used internally in NumPy currently, and also in other places for instance
+    `C++98 <https://en.cppreference.com/w/cpp/memory/allocator/deallocate>`,
+    `C++11 <https://en.cppreference.com/w/cpp/memory/allocator_traits/deallocate>`, and
+    `Rust (allocator_api) <https://doc.rust-lang.org/std/alloc/trait.Allocator.html#tymethod.deallocate>`.
+
+    The consumer of the `PyDataMemAllocator` interface must keep track of ``size`` and make sure it is
+    consistent with the parameter passed to the ``(m|c|re)alloc``  functions.
 
 .. c:function:: const PyDataMem_Handler * PyDataMem_SetHandler(PyDataMem_Handler *handler)
 
@@ -162,8 +173,16 @@ the ``sz`` argument is correct.
     #include <numpy/arrayobject.h>
     NPY_NO_EXPORT void *
 
-    shift_alloc(size_t sz) {
-        char *real = (char *)malloc(sz + 64);
+    typedef struct {
+        void *(*malloc)(size_t);
+        void *(*calloc)(size_t, size_t);
+        void *(*realloc)(void *, size_t);
+        void (*free)(void *);
+    } Allocator;
+
+    NPY_NO_EXPORT void *
+    shift_alloc(Allocator *ctx, size_t sz) {
+        char *real = (char *)ctx->malloc(sz + 64);
         if (real == NULL) {
             return NULL;
         }
@@ -172,8 +191,8 @@ the ``sz`` argument is correct.
     }
 
     NPY_NO_EXPORT void *
-    shift_zero(size_t sz, size_t cnt) {
-        char *real = (char *)calloc(sz + 64, cnt);
+    shift_zero(Allocator *ctx, size_t sz, size_t cnt) {
+        char *real = (char *)ctx->calloc(sz + 64, cnt);
         if (real == NULL) {
             return NULL;
         }
@@ -183,7 +202,7 @@ the ``sz`` argument is correct.
     }
 
     NPY_NO_EXPORT void
-    shift_free(void * p, npy_uintp sz) {
+    shift_free(Allocator *ctx, void * p, npy_uintp sz) {
         if (p == NULL) {
             return ;
         }
@@ -191,37 +210,36 @@ the ``sz`` argument is correct.
         if (strncmp(real, "originally allocated", 20) != 0) {
             fprintf(stdout, "uh-oh, unmatched shift_free, "
                     "no appropriate prefix\\n");
-            /* Make the C runtime crash by calling free on the wrong address */
-            free((char *)p + 10);
-            /* free(real); */
+            /* Make C runtime crash by calling free on the wrong address */
+            ctx->free((char *)p + 10);
+            /* ctx->free(real); */
         }
         else {
-            int i = atoi(real +20);
+            npy_uintp i = (npy_uintp)atoi(real +20);
             if (i != sz) {
-                fprintf(stderr, "uh-oh, unmatched "
-                        "shift_free(ptr, %d) but allocated %d\\n", sz, i);
-                /* Make the C runtime crash by calling free on the wrong address */
-                /* free((char *)p + 10); */
-                free(real);
+                fprintf(stderr, "uh-oh, unmatched shift_free"
+                        "(ptr, %ld) but allocated %ld\\n", sz, i);
+                /* This happens in some places, only print */
+                ctx->free(real);
             }
             else {
-                free(real);
+                ctx->free(real);
             }
         }
     }
 
     NPY_NO_EXPORT void *
-    shift_realloc(void * p, npy_uintp sz) {
+    shift_realloc(Allocator *ctx, void * p, npy_uintp sz) {
         if (p != NULL) {
             char *real = (char *)p - 64;
             if (strncmp(real, "originally allocated", 20) != 0) {
                 fprintf(stdout, "uh-oh, unmatched shift_realloc\\n");
                 return realloc(p, sz);
             }
-            return (void *)((char *)realloc(real, sz + 64) + 64);
+            return (void *)((char *)ctx->realloc(real, sz + 64) + 64);
         }
         else {
-            char *real = (char *)realloc(p, sz + 64);
+            char *real = (char *)ctx->realloc(p, sz + 64);
             if (real == NULL) {
                 return NULL;
             }
@@ -231,63 +249,24 @@ the ``sz`` argument is correct.
         }
     }
 
-    static PyDataMem_Handler new_handler = {
-        "secret_data_allocator",
-        shift_alloc,      /* alloc */
-        shift_zero,       /* zeroed_alloc */
-        shift_free,       /* free */
-        shift_realloc     /* realloc */
+    static Allocator new_handler_ctx = {
+        malloc,
+        calloc,
+        realloc,
+        free
     };
 
-    static PyObject* mem_policy_test_prefix(PyObject *self, PyObject *args)
-    {
-
-        if (!PyArray_Check(args)) {
-            PyErr_SetString(PyExc_ValueError,
-                    "must be called with a numpy scalar or ndarray");
+    static PyDataMem_Handler new_handler = {
+        "secret_data_allocator",
+        {
+            &new_handler_ctx,
+            shift_alloc,      /* malloc */
+            shift_zero, /* calloc */
+            shift_realloc,      /* realloc */
+            shift_free       /* free */
         }
-        return PyUnicode_FromString(
-                        PyDataMem_GetHandlerName((PyArrayObject*)args));
-    };
-
-    static PyObject* mem_policy_set_new_policy(PyObject *self, PyObject *args)
-    {
-
-         const PyDataMem_Handler *old = PyDataMem_SetHandler(&new_handler);
-         return PyUnicode_FromString(old->name);
-
     };
-
-    static PyObject* mem_policy_set_old_policy(PyObject *self, PyObject *args)
-    {
-
-         const PyDataMem_Handler *old = PyDataMem_SetHandler(NULL);
-         return PyUnicode_FromString(old->name);
-
-    };
-
-    static PyMethodDef methods[] = {
-    {"test_prefix", (PyCFunction)mem_policy_test_prefix, METH_O},
-    {"set_new_policy", (PyCFunction)mem_policy_set_new_policy, METH_NOARGS},
-    {"set_old_policy", (PyCFunction)mem_policy_set_old_policy, METH_NOARGS},
-    { NULL }
-    };
-
-    static struct PyModuleDef moduledef = {
-        PyModuleDef_HEAD_INIT,
-        "mem_policy",  /* m_name */
-        NULL,           /* m_doc */
-        -1,             /* m_size */
-        methods,        /* m_methods */
-    };
-
-    PyMODINIT_FUNC
-    PyInit_mem_policy(void) {
-    PyObject *mod = PyModule_Create(&moduledef);
-        import_array();
-        return mod;
-    }
-
+    '''
 
 Related Work
 ------------
@@ -315,7 +294,9 @@ mechanism. The PR was merged with no example code for using these features.
 Discussion
 ----------
 
-Not yet discussed on the mailing list.
+The discussion on the mailing list led to the ``PyDataMemAllocator`` struct
+with a ``context`` field like :c:type:`PyMemAllocatorEx` but with a different
+signature for ``free``.
 
 
 References and Footnotes