baby_mimalloc/
lib.rs

1/*
2 * Copyright 2024 Yufan You
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17//! [Mimalloc](https://github.com/microsoft/mimalloc) implemented in Rust
18//! (not a binding to the C library) with only basic features.
19//!
20//! Lock-free multi-threading, security features, and some performance enhancements are not
21//! implemented.
22//!
23//! It can be used in `no_std` environments.
24//!
25//! # Crate Features
26//!
27//! - **mmap** - Provide [`MimallocMmap`] that uses `mmap` as OS allocator for segments.
28//! - **std_mutex** - Provide [`MimallocMutexWrapper`] that wraps [`Mimalloc`] inside
29//!   [`std::sync::Mutex`] and implements [`GlobalAlloc`].
30//! - **spin_mutex** - Provide [`MimallocMutexWrapper`] that wraps [`Mimalloc`] inside
31//!   [`spin::Mutex`] that can be used in `no_std` environments.
32//! - **deferred_free** - Enable registering a hook to complete deferred free events.
33//!   See the documentation of [`mi_register_deferred_free`](https://microsoft.github.io/mimalloc/group__extended.html#ga3460a6ca91af97be4058f523d3cb8ece).
34
35#![cfg_attr(docsrs, feature(doc_auto_cfg))]
36#![cfg_attr(not(feature = "std"), no_std)]
37
38mod constants;
39mod heap;
40mod list;
41mod page;
42mod segment;
43mod utils;
44
45use core::alloc::{GlobalAlloc, Layout};
46use heap::Heap;
47
48/* wrapper around `heap::Heap` to defined the public API. */
49
50/// The main allocator object.
51///
52/// `A` is the type of the OS allocator for segments.
53///
54/// To use it as the [`global_allocator`], wrap it inside a lock and implement [`GlobalAlloc`].
55/// See [`MimallocMutexWrapper`].
56#[derive(Default)]
57pub struct Mimalloc<A: GlobalAlloc> {
58    heap: Heap,
59    os_alloc: A,
60    #[cfg(feature = "deferred_free")]
61    deferred_free_hook: Option<DeferredFreeHook<A>>,
62}
63
64#[cfg(feature = "deferred_free")]
65pub mod deferred_free;
66#[cfg(feature = "deferred_free")]
67use deferred_free::*;
68
69unsafe impl<A: GlobalAlloc> Send for Mimalloc<A> {}
70
71impl<A: GlobalAlloc> Mimalloc<A> {
72    /// Create a new [`Mimalloc`] instance with an OS allocator.
73    pub const fn with_os_allocator(os_alloc: A) -> Self {
74        Self {
75            heap: Heap::new(),
76            os_alloc,
77            #[cfg(feature = "deferred_free")]
78            deferred_free_hook: None,
79        }
80    }
81
82    #[cfg(feature = "deferred_free")]
83    /// Register a hook to complete deferred free when the allocator needs more memory.
84    /// A new hook replaces the old one.
85    ///
86    /// See the documentation of
87    /// [`mi_register_deferred_free`](https://microsoft.github.io/mimalloc/group__extended.html#ga3460a6ca91af97be4058f523d3cb8ece)
88    /// (the extra `arg` is not supported).
89    pub const fn register_deferred_free(&mut self, hook: DeferredFreeHook<A>) {
90        self.deferred_free_hook = Some(hook);
91    }
92
93    /// Collect free memory.
94    pub fn collect(&mut self) {
95        self.heap.collect(&self.os_alloc);
96    }
97
98    /// [`GlobalAlloc::alloc`] but requires a mutable reference `&mut self`.
99    ///
100    /// # Safety
101    ///
102    /// See [`GlobalAlloc::alloc`].
103    pub unsafe fn alloc(&mut self, layout: Layout) -> *mut u8 {
104        self.heap.malloc_aligned(
105            layout.size(),
106            layout.align(),
107            &self.os_alloc,
108            #[cfg(feature = "deferred_free")]
109            self.deferred_free_hook,
110        )
111    }
112
113    /// [`GlobalAlloc::dealloc`] but requires a mutable reference `&mut self`.
114    ///
115    /// # Safety
116    ///
117    /// See [`GlobalAlloc::dealloc`].
118    pub unsafe fn dealloc(&mut self, ptr: *mut u8, _: Layout) {
119        self.heap.free(ptr, &self.os_alloc)
120    }
121}
122
123impl<A: GlobalAlloc> Drop for Mimalloc<A> {
124    fn drop(&mut self) {
125        self.collect();
126    }
127}
128
129#[cfg(feature = "mmap")]
130mod mmap;
131#[cfg(feature = "mmap")]
132pub use mmap::{new_mimalloc_mmap, MimallocMmap, MmapAlloc};
133
134#[cfg(all(not(docsrs), feature = "std_mutex", feature = "spin_mutex"))]
135compile_error!("Only one of 'std_mutex' and 'spin_mutex' features can be enabled");
136
137#[cfg(any(feature = "std_mutex", feature = "spin_mutex"))]
138mod mutex;
139#[cfg(any(feature = "std_mutex", feature = "spin_mutex"))]
140pub use mutex::MimallocMutexWrapper;
141
142#[cfg(all(feature = "mmap", any(feature = "std_mutex", feature = "spin_mutex")))]
143/// Wrapper around [`Mimalloc`] with `mmap` allocator and mutex.
144pub type MimallocMmapMutex = MimallocMutexWrapper<MmapAlloc>;
145#[cfg(all(feature = "mmap", any(feature = "std_mutex", feature = "spin_mutex")))]
146/// Create a new [`MimallocMmapMutex`] instance by a `const fn`.
147pub const fn new_mimalloc_mmap_mutex() -> MimallocMmapMutex {
148    MimallocMutexWrapper::with_os_allocator(MmapAlloc)
149}
baby_mimalloc/lib.rs

baby_mimalloc/
lib.rs
