fix(alloc): request 16-byte alignment for small bemo_alloc sizes

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::<u8>();
+  // 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::<u8>();
   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::<c_void>()) };
 }
 
@@ -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());

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::<c_void>()
   }
 
+  /// 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;

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;
 }