baby_mimalloc/
mutex.rs

1use crate::Mimalloc;
2use core::alloc::{GlobalAlloc, Layout};
3
4#[cfg(feature = "spin_mutex")]
5use spin::{Mutex, MutexGuard};
6#[cfg(not(feature = "spin_mutex"))]
7use std::sync::{Mutex, MutexGuard};
8
9/// Wrap [`Mimalloc`] inside a [`Mutex`] and implement [`GlobalAlloc`].
10#[derive(Default)]
11pub struct MimallocMutexWrapper<A: GlobalAlloc>(Mutex<Mimalloc<A>>);
12
13impl<A: GlobalAlloc> MimallocMutexWrapper<A> {
14    /// See [`Mimalloc::with_os_allocator`].
15    pub const fn with_os_allocator(os_alloc: A) -> Self {
16        Self(Mutex::new(Mimalloc::with_os_allocator(os_alloc)))
17    }
18
19    #[cfg(feature = "deferred_free")]
20    /// See [`Mimalloc::register_deferred_free`].
21    pub fn register_deferred_free(&self, hook: crate::DeferredFreeHook<A>) {
22        self.allocator().register_deferred_free(hook);
23    }
24
25    /// See [`Mimalloc::collect`].
26    pub fn collect(&self) {
27        self.allocator().collect();
28    }
29
30    fn allocator(&self) -> MutexGuard<Mimalloc<A>> {
31        #[cfg(feature = "spin_mutex")]
32        {
33            self.0.lock()
34        }
35        #[cfg(not(feature = "spin_mutex"))]
36        {
37            self.0.lock().expect("failed to lock the allocator")
38        }
39    }
40}
41
42unsafe impl<A: GlobalAlloc> GlobalAlloc for MimallocMutexWrapper<A> {
43    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
44        self.allocator().alloc(layout)
45    }
46
47    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
48        self.allocator().dealloc(ptr, layout)
49    }
50}