1 #ifndef Py_INTERNAL_OBJECT_ALLOC_H
2 #define Py_INTERNAL_OBJECT_ALLOC_H
3
4 #include "pycore_object.h" // _PyType_HasFeature()
5 #include "pycore_pystate.h" // _PyThreadState_GET()
6 #include "pycore_tstate.h" // _PyThreadStateImpl
7
8 #ifdef __cplusplus
9 extern "C" {
10 #endif
11
12 #ifndef Py_BUILD_CORE
13 # error "this header requires Py_BUILD_CORE define"
14 #endif
15
16 #ifdef Py_GIL_DISABLED
17 static inline mi_heap_t *
_PyObject_GetAllocationHeap(_PyThreadStateImpl * tstate,PyTypeObject * tp)18 _PyObject_GetAllocationHeap(_PyThreadStateImpl *tstate, PyTypeObject *tp)
19 {
20 struct _mimalloc_thread_state *m = &tstate->mimalloc;
21 if (_PyType_HasFeature(tp, Py_TPFLAGS_PREHEADER)) {
22 return &m->heaps[_Py_MIMALLOC_HEAP_GC_PRE];
23 }
24 else if (_PyType_IS_GC(tp)) {
25 return &m->heaps[_Py_MIMALLOC_HEAP_GC];
26 }
27 else {
28 return &m->heaps[_Py_MIMALLOC_HEAP_OBJECT];
29 }
30 }
31 #endif
32
33 // Sets the heap used for PyObject_Malloc(), PyObject_Realloc(), etc. calls in
34 // Py_GIL_DISABLED builds. We use different heaps depending on if the object
35 // supports GC and if it has a pre-header. We smuggle the choice of heap
36 // through the _mimalloc_thread_state. In the default build, this simply
37 // calls PyObject_Malloc().
38 static inline void *
_PyObject_MallocWithType(PyTypeObject * tp,size_t size)39 _PyObject_MallocWithType(PyTypeObject *tp, size_t size)
40 {
41 #ifdef Py_GIL_DISABLED
42 _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET();
43 struct _mimalloc_thread_state *m = &tstate->mimalloc;
44 m->current_object_heap = _PyObject_GetAllocationHeap(tstate, tp);
45 #endif
46 void *mem = PyObject_Malloc(size);
47 #ifdef Py_GIL_DISABLED
48 m->current_object_heap = &m->heaps[_Py_MIMALLOC_HEAP_OBJECT];
49 #endif
50 return mem;
51 }
52
53 static inline void *
_PyObject_ReallocWithType(PyTypeObject * tp,void * ptr,size_t size)54 _PyObject_ReallocWithType(PyTypeObject *tp, void *ptr, size_t size)
55 {
56 #ifdef Py_GIL_DISABLED
57 _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET();
58 struct _mimalloc_thread_state *m = &tstate->mimalloc;
59 m->current_object_heap = _PyObject_GetAllocationHeap(tstate, tp);
60 #endif
61 void *mem = PyObject_Realloc(ptr, size);
62 #ifdef Py_GIL_DISABLED
63 m->current_object_heap = &m->heaps[_Py_MIMALLOC_HEAP_OBJECT];
64 #endif
65 return mem;
66 }
67
68 #ifdef __cplusplus
69 }
70 #endif
71 #endif // !Py_INTERNAL_OBJECT_ALLOC_H
72