fix: chunk allocation wrt size alignment for downwards bumping (fixes #…

…32)

src/chunk_raw.rs

@@ -9,7 +9,7 @@ use crate::{
     ChunkSize, ErrorBehavior, MinimumAlignment, SizedTypeProperties, SupportedMinimumAlignment, CHUNK_ALIGN_MIN,
 };
 
-use allocator_api2::alloc::Allocator;
+use allocator_api2::alloc::{AllocError, Allocator};
 
 /// Represents an allocated chunk.
 ///
@@ -54,16 +54,37 @@ impl<const UP: bool, A> RawChunk<UP, A> {
         A: Allocator,
         for<'a> &'a A: Allocator,
     {
-        let ptr = size.allocate(&allocator)?;
+        let layout = size.layout();
 
-        let size = ptr.len();
-        let ptr = ptr.cast::<u8>();
+        let allocation = match allocator.allocate(layout) {
+            Ok(ok) => ok,
+            Err(AllocError) => return Err(E::allocation(layout)),
+        };
+
+        // The size must always be a multiple of `CHUNK_ALIGN_MIN`.
+        // We use optimizations in `alloc` that make use of this.
+        //
+        // If `!UP`, the size must also be an aligned to `ChunkHeader<_>`
+        // so the header can live at the end.
+        let downwards_align = if UP {
+            CHUNK_ALIGN_MIN
+        } else {
+            CHUNK_ALIGN_MIN.max(align_of::<ChunkHeader<A>>())
+        };
 
-        // The chunk size must always be a multiple of `CHUNK_ALIGN_MIN`.
-        // We use optimizations in `alloc` that require this.
-        // `ChunkSize::allocate` has already trimmed the allocation size to a multiple of `CHUNK_ALIGN_MIN`
+        // This truncation can not result in a size smaller than the layout's size because
+        // we truncated the layout's size in the same way (see ChunkSize::layout).
+        //
+        // NB: The size must not be smaller than layout's size, so it still [fits] the
+        // memory block so we can deallocate with that size.
+        //
+        // [fits]: https://doc.rust-lang.org/std/alloc/trait.Allocator.html#memory-fitting
+        let size = down_align_usize(allocation.len(), downwards_align);
+        debug_assert!(size >= layout.size());
         debug_assert!(size % CHUNK_ALIGN_MIN == 0);
 
+        let ptr = allocation.cast::<u8>();
+
         let prev = prev.map(|c| c.header);
         let next = Cell::new(None);
 
@@ -81,9 +102,6 @@ impl<const UP: bool, A> RawChunk<UP, A> {
 
                 header
             } else {
-                // `Chunk::allocate` allocated a size that is a multiple of `align_of::<ChunkHeader<A>>()`.
-                debug_assert!(size % align_of::<ChunkHeader<A>>() == 0);
-
                 let header = nonnull::sub(nonnull::add(ptr, size).cast::<ChunkHeader<A>>(), 1);
 
                 header.as_ptr().write(ChunkHeader {

src/chunk_size.rs

@@ -4,16 +4,14 @@ use core::{
     marker::PhantomData,
     mem::{align_of, size_of},
     num::NonZeroUsize,
-    ptr::NonNull,
 };
 
 use crate::{
-    polyfill::{const_unwrap, nonnull, nonzero},
-    up_align_nonzero, ChunkHeader, ErrorBehavior, CHUNK_ALIGN_MIN,
+    down_align_usize,
+    polyfill::{const_unwrap, nonzero},
+    up_align_nonzero, ChunkHeader, CHUNK_ALIGN_MIN,
 };
 
-use allocator_api2::alloc::Allocator;
-
 /// We leave some space per allocation for the base allocator.
 pub(crate) type AssumedMallocOverhead = [*const u8; 2];
 pub(crate) const ASSUMED_PAGE_SIZE: NonZeroUsize = const_unwrap(NonZeroUsize::new(0x1000));
@@ -113,7 +111,15 @@ impl<const UP: bool, A> ChunkSize<UP, A> {
     pub(crate) fn layout(self) -> Layout {
         // we checked in `new` that we can create a layout from this size
 
-        let size_for_layout = self.0.get() - size_of::<AssumedMallocOverhead>();
+        let size_without_overhead = self.0.get() - size_of::<AssumedMallocOverhead>();
+
+        let downwards_align = if UP {
+            CHUNK_ALIGN_MIN
+        } else {
+            CHUNK_ALIGN_MIN.max(align_of::<ChunkHeader<A>>())
+        };
+
+        let size_for_layout = down_align_usize(size_without_overhead, downwards_align);
         let align = align_of::<ChunkHeader<A>>();
 
         unsafe { Layout::from_size_align_unchecked(size_for_layout, align) }
@@ -127,33 +133,6 @@ impl<const UP: bool, A> ChunkSize<UP, A> {
             other
         }
     }
-
-    #[inline]
-    pub(crate) fn allocate<E: ErrorBehavior>(self, allocator: impl Allocator) -> Result<NonNull<[u8]>, E> {
-        let layout = self.layout();
-
-        let slice = match allocator.allocate(layout) {
-            Ok(slice) => slice,
-            Err(_) => return Err(E::allocation(layout)),
-        };
-
-        let size = slice.len();
-        let ptr = slice.cast::<u8>();
-
-        // `ptr + size` must be an aligned to `CHUNK_ALIGN_MIN`
-        // if `!UP`, `ptr + size` must also be an aligned `*const ChunkHeader<_>`
-        let down_alignment = if UP {
-            CHUNK_ALIGN_MIN
-        } else {
-            CHUNK_ALIGN_MIN.max(align_of::<ChunkHeader<A>>())
-        };
-
-        let truncated_size = down_align(size, down_alignment);
-        debug_assert!(truncated_size >= layout.size());
-
-        let truncated_slice = nonnull::slice_from_raw_parts(ptr, truncated_size);
-        Ok(truncated_slice)
-    }
 }
 
 #[derive(Clone, Copy)]
@@ -211,10 +190,3 @@ const fn up_align(addr: usize, align: usize) -> Option<usize> {
     let aligned = addr_plus_mask & !mask;
     Some(aligned)
 }
-
-#[inline(always)]
-fn down_align(addr: usize, align: usize) -> usize {
-    debug_assert!(align.is_power_of_two());
-    let mask = align - 1;
-    addr & !mask
-}

src/tests/chunk_size.rs

@@ -21,7 +21,7 @@ fn aligned_allocator_issue_32() {
         }
 
         unsafe fn deallocate(&self, ptr: core::ptr::NonNull<u8>, layout: core::alloc::Layout) {
-            Global.deallocate(ptr, layout)
+            Global.deallocate(ptr, layout);
         }
     }