diff --git a/src/internal.rs b/src/internal.rs index d65bbdf..2fad5e2 100644 --- a/src/internal.rs +++ b/src/internal.rs @@ -74,14 +74,13 @@ pub(crate) struct Large(pub(crate) NonNull>); impl Large { pub(crate) fn new_layout(cap: usize) -> Layout { - let additional_size = size_of::().checked_mul(cap).expect("new capacity is too large"); - let size = size_of::>() - .checked_add(additional_size) - .expect("new capacity is too large"); - let align = align_of::>(); - // SAFETY: size of Allocated must be a multiple of align of Allocated, - // which must be a multiple of align of T due to the `data` field. - unsafe { Layout::from_size_align_unchecked(size, align) } + let header = Layout::new::>(); + let data = Layout::array::(cap).expect("new capacity is too large"); + let (layout, offset) = header.extend(data).expect("new capacity is too large"); + // `Layout::new::>()` already includes padding for `T`, + // so the trailing data starts immediately after the header. + debug_assert_eq!(offset, size_of::>()); + layout.pad_to_align() } pub(crate) fn as_allocated(&self) -> (&Allocated, *const T) { @@ -262,6 +261,9 @@ impl Allocated { /// /// The data behind the header are allowed to be uninitialized. pub(crate) unsafe fn data_start(this: NonNull) -> *mut T { - unsafe { (&raw mut (*this.as_ptr()).data_start).cast() } + let base = this.as_ptr().cast::(); + // SAFETY: base points to a valid allocation whose data starts immediately after the header. + // `new_layout` uses the same header size (see the debug assertion there). + unsafe { base.add(size_of::()).cast::() } } } diff --git a/src/tests.rs b/src/tests.rs index fc5ba16..a0a64d9 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -1,5 +1,6 @@ use alloc::string::{String, ToString}; use alloc::vec::Vec; +use alloc::alloc::dealloc; use core::cell::Cell; use core::panic::AssertUnwindSafe; use core::{mem, ops}; @@ -14,6 +15,49 @@ fn assert_size() { assert_eq!(size_of::>(), 8); } +#[test] +fn test_large_layout_padding_and_data_start() { + use crate::internal::{Allocated, Large}; + + #[allow(dead_code)] + #[repr(align(32))] + struct Align32(u8); + + fn assert_layout() { + use core::alloc::Layout; + use core::ptr::NonNull; + + let layout = Large::::new_layout(1); + assert_eq!(layout.size() % layout.align(), 0); + + let large = Large::::new_empty(1); + let layout = large.current_layout(); + assert_eq!(layout, Large::::new_layout(1)); + struct DeallocGuard { + ptr: NonNull>, + layout: Layout, + } + impl Drop for DeallocGuard { + fn drop(&mut self) { + unsafe { + dealloc(self.ptr.as_ptr().cast(), self.layout); + } + } + } + let _guard = DeallocGuard { ptr: large.0, layout }; + + let (allocated, data_start) = large.as_allocated(); + let base = allocated as *const Allocated as *const u8; + let offset = unsafe { data_start.cast::().offset_from(base) as usize }; + assert_eq!(offset, size_of::>()); + } + + assert_layout::(); + assert_layout::(); + assert_layout::(); + assert_layout::(); +} + struct AssertDrop<'a> { string: String, counter: &'a Cell,