mimalloc: enable as allocator option

DinoV · DinoV · commit f378513e61fe · 2023-09-28T12:46:22.000-07:00
diff --git a/Include/cpython/pymem.h b/Include/cpython/pymem.h
@@ -29,6 +29,10 @@ typedef enum {
     PYMEM_ALLOCATOR_PYMALLOC = 5,
     PYMEM_ALLOCATOR_PYMALLOC_DEBUG = 6,
 #endif
+#ifdef WITH_MIMALLOC
+    PYMEM_ALLOCATOR_MIMALLOC = 7,
+    PYMEM_ALLOCATOR_MIMALLOC_DEBUG = 8,
+#endif
 } PyMemAllocatorName;
 
 
diff --git a/Include/internal/pycore_pymem_init.h b/Include/internal/pycore_pymem_init.h
@@ -18,17 +18,31 @@ extern void * _PyMem_RawRealloc(void *, void *, size_t);
 extern void _PyMem_RawFree(void *, void *);
 #define PYRAW_ALLOC {NULL, _PyMem_RawMalloc, _PyMem_RawCalloc, _PyMem_RawRealloc, _PyMem_RawFree}
 
-#ifdef WITH_PYMALLOC
+#ifdef WITH_MIMALLOC
+extern void *_PyMem_MiMalloc(void *ctx, size_t size);
+extern void * _PyMem_MiCalloc(void *ctx, size_t nelem, size_t elsize);
+extern void *_PyMem_MiRealloc(void *ctx, void *ptr, size_t size);
+extern void _PyMem_MiFree(void *ctx, void *ptr);
+extern void * _PyObject_MiMalloc(void *ctx, size_t nbytes);
+extern void* _PyObject_MiCalloc(void *ctx, size_t nelem, size_t elsize);
+extern void* _PyObject_MiRealloc(void *ctx, void *ptr, size_t nbytes);
+extern void _PyObject_MiFree(void *ctx, void *ptr);
+
+#  define PYMEM_ALLOC {NULL, _PyMem_MiMalloc, _PyMem_MiCalloc, _PyMem_MiRealloc, _PyMem_MiFree}
+#  define PYOBJ_ALLOC {NULL, _PyObject_MiMalloc, _PyObject_MiCalloc, _PyObject_MiRealloc, _PyObject_MiFree}
+
+#elif defined(WITH_PYMALLOC)
 extern void* _PyObject_Malloc(void *, size_t);
 extern void* _PyObject_Calloc(void *, size_t, size_t);
 extern void _PyObject_Free(void *, void *);
 extern void* _PyObject_Realloc(void *, void *, size_t);
 #  define PYOBJ_ALLOC {NULL, _PyObject_Malloc, _PyObject_Calloc, _PyObject_Realloc, _PyObject_Free}
+#  define PYMEM_ALLOC PYOBJ_ALLOC
 #else
 # define PYOBJ_ALLOC PYRAW_ALLOC
+# define PYMEM_ALLOC PYOBJ_ALLOC
 #endif  // WITH_PYMALLOC
 
-#define PYMEM_ALLOC PYOBJ_ALLOC
 
 extern void* _PyMem_DebugRawMalloc(void *, size_t);
 extern void* _PyMem_DebugRawCalloc(void *, size_t, size_t);
diff --git a/Lib/test/support/__init__.py b/Lib/test/support/__init__.py
@@ -2028,6 +2028,11 @@ def with_pymalloc():
     return _testcapi.WITH_PYMALLOC
 
 
+def with_mimalloc():
+    import _testcapi
+    return _testcapi.WITH_MIMALLOC
+
+
 class _ALWAYS_EQ:
     """
     Object that is equal to anything.
diff --git a/Lib/test/test_cmd_line.py b/Lib/test/test_cmd_line.py
@@ -724,7 +724,9 @@ def test_xdev(self):
             code = "import _testinternalcapi; print(_testinternalcapi.pymem_getallocatorsname())"
             with support.SuppressCrashReport():
                 out = self.run_xdev("-c", code, check_exitcode=False)
-            if support.with_pymalloc():
+            if support.with_mimalloc():
+                alloc_name = "mimalloc_debug"
+            elif support.with_pymalloc():
                 alloc_name = "pymalloc_debug"
             else:
                 alloc_name = "malloc_debug"
@@ -803,7 +805,11 @@ def check_pythonmalloc(self, env_var, name):
     def test_pythonmalloc(self):
         # Test the PYTHONMALLOC environment variable
         pymalloc = support.with_pymalloc()
-        if pymalloc:
+        mimalloc = support.with_mimalloc()
+        if mimalloc:
+            default_name = 'mimalloc_debug' if support.Py_DEBUG else 'mimalloc'
+            default_name_debug = 'mimalloc_debug'
+        elif pymalloc:
             default_name = 'pymalloc_debug' if support.Py_DEBUG else 'pymalloc'
             default_name_debug = 'pymalloc_debug'
         else:
@@ -821,6 +827,11 @@ def test_pythonmalloc(self):
                 ('pymalloc', 'pymalloc'),
                 ('pymalloc_debug', 'pymalloc_debug'),
             ))
+        if mimalloc:
+            tests.extend((
+                ('mimalloc', 'mimalloc'),
+                ('mimalloc_debug', 'mimalloc_debug'),
+            ))
 
         for env_var, name in tests:
             with self.subTest(env_var=env_var, name=name):
diff --git a/Modules/_testcapi/mem.c b/Modules/_testcapi/mem.c
@@ -613,5 +613,14 @@ _PyTestCapi_Init_Mem(PyObject *mod)
         return -1;
     }
 
+#ifdef WITH_MIMALLOC
+    v = Py_True;
+#else
+    v = Py_False;
+#endif
+    if (PyModule_AddObjectRef(mod, "WITH_MIMALLOC", v) < 0) {
+        return -1;
+    }
+
     return 0;
 }
diff --git a/Objects/obmalloc.c b/Objects/obmalloc.c

Original file line number	Diff line number	Diff line change
`@@ -613,5 +613,14 @@ _PyTestCapi_Init_Mem(PyObject *mod)`
`613`	`613`	`return -1;`
`614`	`614`	`}`
`615`	`615`
	`616`	`+#ifdef WITH_MIMALLOC`
	`617`	`+ v = Py_True;`
	`618`	`+#else`
	`619`	`+ v = Py_False;`
	`620`	`+#endif`
	`621`	`+ if (PyModule_AddObjectRef(mod, "WITH_MIMALLOC", v) < 0) {`
	`622`	`+ return -1;`
	`623`	`+ }`
	`624`	`+`
`616`	`625`	`return 0;`
`617`	`626`	`}`