Skip to main content

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}