diff --git a/crates/bemo/src/alloc.rs b/crates/bemo/src/alloc.rs index f48690f..ed0bb05 100644 --- a/crates/bemo/src/alloc.rs +++ b/crates/bemo/src/alloc.rs @@ -11,7 +11,7 @@ use std::ffi::c_void; use std::sync::atomic::{AtomicU64, Ordering}; -use crate::buffer::{mi_free, mi_malloc}; +use crate::buffer::{mi_free, mi_malloc_aligned}; /// Largest single allocation, matching the largest Java direct buffer. pub const MAX_ALLOCATION: usize = i32::MAX as usize; @@ -24,8 +24,8 @@ pub fn allocate(size: usize) -> *mut u8 { if size == 0 || size > MAX_ALLOCATION { return std::ptr::null_mut(); } - // SAFETY: size is nonzero and bounded well below isize::MAX. - let pointer = unsafe { mi_malloc(size) }.cast::(); + // SAFETY: size is nonzero and bounded well below isize::MAX; pair with free at this size. + let pointer = unsafe { mi_malloc_aligned(size, 16) }.cast::(); if !pointer.is_null() { LIVE.fetch_add(size as u64, Ordering::Relaxed); } @@ -41,7 +41,7 @@ pub unsafe fn free(pointer: *mut u8, size: usize) { return; } LIVE.fetch_sub(size as u64, Ordering::Relaxed); - // SAFETY: The caller guarantees pointer came from mi_malloc through allocate and is live. + // SAFETY: The caller guarantees pointer came from mi_malloc_aligned through allocate and is live. unsafe { mi_free(pointer.cast::()) }; } @@ -69,6 +69,41 @@ mod tests { unsafe { free(pointer, 4096) }; } + /// Sizes 1-8 are the range mimalloc serves from its 8-byte block-size class, where a plain + /// `mi_malloc` returns roughly half of pointers congruent to 8 mod 16. Hold many allocations of + /// each size at once so mimalloc must hand out sequential blocks from the same small class + /// rather than masking the residue with one freshly 16-aligned block; every pointer must be + /// 16-aligned and every byte writable. Sizes 9-64 confirm the larger 16-multiple classes are + /// not regressed. + #[test] + fn every_size_class_returns_16_aligned_storage() { + for size in 1usize..=64 { + let mut held: Vec<*mut u8> = Vec::with_capacity(32); + for _ in 0..32 { + let pointer = allocate(size); + assert!(!pointer.is_null(), "size {size} should allocate"); + assert_eq!( + pointer as usize % 16, + 0, + "size {size} returned a pointer not aligned to 16" + ); + // SAFETY: allocate returned `size` writable bytes; fill all of them. + unsafe { pointer.write_bytes(0xa5 ^ size as u8, size) }; + // SAFETY: The first byte is within the allocation and was just initialized. + let first = unsafe { pointer.read() }; + // SAFETY: The final byte is within the allocation and was just initialized. + let last = unsafe { pointer.add(size - 1).read() }; + assert_eq!(first, 0xa5 ^ size as u8, "size {size} first byte"); + assert_eq!(last, 0xa5 ^ size as u8, "size {size} last byte"); + held.push(pointer); + } + for pointer in held { + // SAFETY: Every held pointer came from allocate(size) and is freed once with that size. + unsafe { free(pointer, size) }; + } + } + } + #[test] fn invalid_sizes_and_null_frees_do_nothing() { assert!(allocate(0).is_null()); diff --git a/crates/bemo/src/buffer.rs b/crates/bemo/src/buffer.rs index 7e7d901..71322ce 100644 --- a/crates/bemo/src/buffer.rs +++ b/crates/bemo/src/buffer.rs @@ -33,7 +33,7 @@ compile_error!("Select either bundled-mimalloc or rust-allocator, not both"); not(miri), not(target_vendor = "apple") ))] -pub(crate) use libmimalloc_sys::{mi_free, mi_malloc}; +pub(crate) use libmimalloc_sys::{mi_free, mi_malloc, mi_malloc_aligned}; #[cfg(all( not(feature = "bundled-mimalloc"), @@ -43,11 +43,14 @@ pub(crate) use libmimalloc_sys::{mi_free, mi_malloc}; ))] unsafe extern "C" { pub(crate) fn mi_malloc(size: usize) -> *mut c_void; + /// Aligned counterpart to [`mi_malloc`]; freed with the same [`mi_free`]. The embedder's + /// mimalloc supplies this alongside `mi_malloc`/`mi_free` (see `docs/elide-integration.md`). + pub(crate) fn mi_malloc_aligned(size: usize, alignment: usize) -> *mut c_void; pub(crate) fn mi_free(pointer: *mut c_void); } #[cfg(any(miri, target_vendor = "apple", feature = "rust-allocator"))] -pub(crate) use rust_alloc::{mi_free, mi_malloc}; +pub(crate) use rust_alloc::{mi_free, mi_malloc, mi_malloc_aligned}; // The bundled allocator is still linked under miri, just never called from here. #[cfg(all(feature = "bundled-mimalloc", miri, not(target_vendor = "apple")))] @@ -82,6 +85,24 @@ mod rust_alloc { unsafe { base.add(MI_ALIGN) }.cast::() } + /// Aligned counterpart to [`mi_malloc`], pairing with the same [`mi_free`]. [`mi_malloc`] already + /// lays out every payload at exactly `MI_ALIGN`, so any requested alignment up to `MI_ALIGN` is + /// satisfied by the same header-backed allocation. The transport only ever requests `MI_ALIGN`, + /// which is all the fixed `MI_ALIGN` header and [`mi_free`] can reconstruct; a larger alignment + /// would need a wider offset that [`mi_free`] does not read. + pub(crate) unsafe fn mi_malloc_aligned(size: usize, alignment: usize) -> *mut c_void { + // SAFETY: alignment is a power of two no greater than MI_ALIGN, so mi_malloc's MI_ALIGN-aligned + // payload satisfies it and the fixed header keeps pairing with mi_free. + debug_assert!( + alignment.is_power_of_two() && alignment <= MI_ALIGN, + "unsupported alignment {alignment}" + ); + // SAFETY: mi_malloc returns an MI_ALIGN-aligned payload (satisfying any + // power-of-two alignment up to MI_ALIGN) with the fixed MI_ALIGN header + // mi_free reconstructs, or null on allocation failure. + unsafe { mi_malloc(size) } + } + pub(crate) unsafe fn mi_free(pointer: *mut c_void) { if pointer.is_null() { return; diff --git a/tests/static.c b/tests/static.c index 72551af..cfee542 100644 --- a/tests/static.c +++ b/tests/static.c @@ -25,5 +25,16 @@ int main(void) { assert(bemo_alloc(UINT64_C(1) << 31) == 0); bemo_free(0, 4096); assert(bemo_alloc_live() == live); + for (uint64_t size = 1; size <= 16; size++) { + unsigned char *small = bemo_alloc(size); + assert(small != 0 && ((uintptr_t)small & 15) == 0); + assert(bemo_alloc_live() == live + size); + for (uint64_t i = 0; i < size; i++) + small[i] = (unsigned char)(size + i); + assert(small[0] == (unsigned char)size); + assert(small[size - 1] == (unsigned char)(size + (size - 1))); + bemo_free(small, size); + assert(bemo_alloc_live() == live); + } return 0; }