moirai_core/memory/allocator.rs
1use crate::memory::CACHE_LINE_SIZE;
2use std::alloc::Layout;
3use std::ptr::NonNull;
4
5/// Cache-aligned memory allocator for high-performance data structures.
6pub struct CacheAlignedAllocator;
7
8/// Layout of `count` values of `T` starting on a cache-line boundary.
9///
10/// Transfer granularity: this places the *start* of the array on a line
11/// boundary so element 0 does not straddle two lines. Separating two
12/// concurrently written atomics is a different problem, solved by
13/// `CacheAligned` at the field level, not by widening this alignment. The size
14/// comes from `Layout::array`, so a count whose byte size overflows `isize` is
15/// refused instead of wrapping into an undersized allocation.
16fn cache_aligned_array<T>(count: usize) -> Option<Layout> {
17 Layout::array::<T>(count)
18 .ok()?
19 .align_to(CACHE_LINE_SIZE)
20 .ok()
21}
22
23impl CacheAlignedAllocator {
24 /// Allocate cache-aligned memory for optimal performance
25 pub fn allocate<T>(count: usize) -> Option<NonNull<T>> {
26 let layout = cache_aligned_array::<T>(count)?;
27 // A zero-sized layout violates `GlobalAlloc::alloc`'s contract.
28 if layout.size() == 0 {
29 return None;
30 }
31
32 // SAFETY: `layout` is valid and non-zero-sized; allocation failure
33 // is surfaced as `None` through the null check.
34 unsafe {
35 #[cfg(feature = "mnemosyne")]
36 {
37 use core::alloc::GlobalAlloc;
38 let ptr = mnemosyne::Mnemosyne.alloc(layout);
39 NonNull::new(ptr.cast::<T>())
40 }
41 #[cfg(not(feature = "mnemosyne"))]
42 {
43 let ptr = std::alloc::alloc(layout);
44 NonNull::new(ptr.cast::<T>())
45 }
46 }
47 }
48
49 /// Deallocate cache-aligned memory
50 ///
51 /// # Safety
52 ///
53 /// The caller must ensure that:
54 /// - `ptr` was allocated by `allocate` with the same type and count
55 /// - `ptr` is valid and properly aligned
56 /// - No other references to the memory exist
57 /// - The memory is not accessed after deallocation
58 pub unsafe fn deallocate<T>(ptr: NonNull<T>, count: usize) {
59 unsafe {
60 if let Some(layout) = cache_aligned_array::<T>(count) {
61 #[cfg(feature = "mnemosyne")]
62 {
63 use core::alloc::GlobalAlloc;
64 mnemosyne::Mnemosyne.dealloc(ptr.as_ptr().cast::<u8>(), layout);
65 }
66 #[cfg(not(feature = "mnemosyne"))]
67 {
68 std::alloc::dealloc(ptr.as_ptr().cast::<u8>(), layout);
69 }
70 }
71 }
72 }
73}