regorus_mimalloc/
mimalloc.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4use core::alloc::{GlobalAlloc, Layout};
5use core::ffi::c_void;
6use mimalloc_sys::{
7    mi_free, mi_malloc_aligned, mi_realloc_aligned, mi_zalloc_aligned, MI_ALIGNMENT_MAX,
8};
9
10#[cfg(feature = "allocator-memory-limits")]
11pub use crate::limits::{
12    allocation_stats_snapshot, current_thread_allocation_stats, flush_thread_counters,
13    global_allocation_stats_snapshot, record_alloc, record_free, set_thread_flush_threshold,
14    GlobalAllocationStats, ThreadAllocationStats,
15};
16pub struct Mimalloc;
17
18unsafe impl GlobalAlloc for Mimalloc {
19    #[inline]
20    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
21        debug_assert!(layout.align() < MI_ALIGNMENT_MAX);
22        let ptr = mi_malloc_aligned(layout.size(), layout.align()).cast::<u8>();
23
24        #[cfg(feature = "allocator-memory-limits")]
25        if !ptr.is_null() {
26            record_alloc(layout.size());
27        }
28        ptr
29    }
30
31    #[inline]
32    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
33        #[cfg(feature = "allocator-memory-limits")]
34        record_free(_layout.size());
35        mi_free(ptr.cast::<c_void>());
36    }
37
38    #[inline]
39    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
40        debug_assert!(layout.align() < MI_ALIGNMENT_MAX);
41        let ptr = mi_zalloc_aligned(layout.size(), layout.align()).cast::<u8>();
42
43        #[cfg(feature = "allocator-memory-limits")]
44        if !ptr.is_null() {
45            record_alloc(layout.size());
46        }
47        ptr
48    }
49
50    #[inline]
51    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
52        debug_assert!(layout.align() < MI_ALIGNMENT_MAX);
53        let result =
54            mi_realloc_aligned(ptr.cast::<c_void>(), new_size, layout.align()).cast::<u8>();
55
56        #[cfg(feature = "allocator-memory-limits")]
57        if !result.is_null() {
58            record_free(layout.size());
59            record_alloc(new_size);
60        }
61        result
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68    use std::error::Error;
69
70    #[test]
71    fn memory_can_be_allocated_and_freed() -> Result<(), Box<dyn Error>> {
72        let layout = Layout::from_size_align(8, 8)?;
73        let alloc = Mimalloc;
74
75        unsafe {
76            let ptr = alloc.alloc(layout);
77            assert!(!ptr.cast::<c_void>().is_null());
78            alloc.dealloc(ptr, layout);
79        }
80        Ok(())
81    }
82
83    #[test]
84    fn memory_can_be_alloc_zeroed_and_freed() -> Result<(), Box<dyn Error>> {
85        let layout = Layout::from_size_align(8, 8)?;
86        let alloc = Mimalloc;
87
88        unsafe {
89            let ptr = alloc.alloc_zeroed(layout);
90            assert!(!ptr.cast::<c_void>().is_null());
91            alloc.dealloc(ptr, layout);
92        }
93        Ok(())
94    }
95
96    #[test]
97    fn large_chunks_of_memory_can_be_allocated_and_freed() -> Result<(), Box<dyn Error>> {
98        let layout = Layout::from_size_align(2 * 1024 * 1024 * 1024, 8)?;
99        let alloc = Mimalloc;
100
101        unsafe {
102            let ptr = alloc.alloc(layout);
103            assert!(!ptr.cast::<c_void>().is_null());
104            alloc.dealloc(ptr, layout);
105        }
106        Ok(())
107    }
108
109    #[cfg(feature = "allocator-memory-limits")]
110    #[test]
111    fn allocation_stats_report_live_usage() -> Result<(), Box<dyn Error>> {
112        let layout = Layout::from_size_align(4 * 1024, 8)?;
113        let alloc = Mimalloc;
114
115        let before_global = global_allocation_stats_snapshot().allocated;
116        let before_thread = current_thread_allocation_stats().allocated;
117
118        unsafe {
119            let ptr = alloc.alloc(layout);
120            assert!(!ptr.cast::<c_void>().is_null());
121            let during_thread = current_thread_allocation_stats().allocated;
122            assert!(during_thread >= before_thread + layout.size() as i64);
123            alloc.dealloc(ptr, layout);
124        }
125
126        let after_thread = current_thread_allocation_stats().allocated;
127        assert_eq!(after_thread, before_thread);
128
129        let after_global = global_allocation_stats_snapshot().allocated;
130        assert!(after_global >= before_global);
131        Ok(())
132    }
133}
regorus_mimalloc/mimalloc.rs

regorus_mimalloc/
mimalloc.rs
