Add StructAlloc

StructAlloc<T> is an internal allocator that, when asked to allocate
memory for a type `Data`, actually allocates space for

```rust
struct Struct {
    t: T,
    data: Data,
}
```

This is useful for allocations that will later be stored in smart
pointers. Instead of having to copy data, we simply offset the pointer,
knowing that we can store the smart metadata (e.g. ref count) in the
space before the data allocation.

Author: Julian Orth

library/alloc/src/alloc.rs

@@ -16,6 +16,8 @@ pub use core::alloc::*;
 
 #[cfg(test)]
 mod tests;
+#[macro_use]
+pub(crate) mod struct_alloc;
 
 extern "Rust" {
     // These are the magic symbols to call the global allocator.  rustc generates

library/alloc/src/alloc/struct_alloc.rs

@@ -0,0 +1,221 @@
+use crate::alloc::Global;
+use crate::fmt;
+use core::alloc::{AllocError, Allocator, Layout};
+use core::fmt::{Debug, Formatter};
+use core::marker::PhantomData;
+use core::mem;
+use core::ptr::NonNull;
+
+#[cfg(test)]
+mod tests;
+
+/// Allocator that adds appropriate padding for a repr(C) struct.
+///
+/// This allocator takes as type arguments the type of a field `T` and an allocator `A`.
+///
+/// Consider
+///
+/// ```rust,ignore (not real code)
+/// #[repr(C)]
+/// struct Struct {
+///     t: T,
+///     data: Data,
+/// }
+/// ```
+///
+/// where `Data` is a type with layout `layout`.
+///
+/// When this allocator creates an allocation for layout `layout`, the pointer can be
+/// offset by `-offsetof(Struct, data)` and the resulting pointer points is an allocation
+/// of `A` for `Layout::new::<Struct>()`.
+pub(crate) struct StructAlloc<T, A = Global>(A, PhantomData<*const T>);
+
+impl<T, A> StructAlloc<T, A> {
+    #[allow(dead_code)]
+    /// Creates a new allocator.
+    pub(crate) fn new(allocator: A) -> Self {
+        Self(allocator, PhantomData)
+    }
+
+    /// Computes the layout of `Struct`.
+    fn struct_layout(data_layout: Layout) -> Result<Layout, AllocError> {
+        let t_align = mem::align_of::<T>();
+        let t_size = mem::size_of::<T>();
+        if t_size == 0 && t_align == 1 {
+            // Skip the checks below
+            return Ok(data_layout);
+        }
+        let data_align = data_layout.align();
+        // The contract of `Layout` guarantees that `data_align > 0`.
+        let data_align_minus_1 = data_align.wrapping_sub(1);
+        let data_size = data_layout.size();
+        let align = data_align.max(t_align);
+        let align_minus_1 = align.wrapping_sub(1);
+        // `size` is
+        //     t_size rounded up to `data_align`
+        // plus
+        //     `data_size` rounded up to `align`
+        // Note that the result is a multiple of `align`.
+        let (t_size_aligned, t_overflow) =
+            t_size.overflowing_add(t_size.wrapping_neg() & data_align_minus_1);
+        let (data_size_aligned, data_overflow) = match data_size.overflowing_add(align_minus_1) {
+            (sum, req_overflow) => (sum & !align_minus_1, req_overflow),
+        };
+        let (size, sum_overflow) = t_size_aligned.overflowing_add(data_size_aligned);
+        if t_overflow || data_overflow || sum_overflow {
+            return Err(AllocError);
+        }
+        unsafe { Ok(Layout::from_size_align_unchecked(size, align)) }
+    }
+
+    /// Returns the offset of `data` in `Struct`.
+    #[inline]
+    pub(crate) fn offset_of_data(data_layout: Layout) -> usize {
+        let t_size = mem::size_of::<T>();
+        // The contract of `Layout` guarantees `.align() > 0`
+        let data_align_minus_1 = data_layout.align().wrapping_sub(1);
+        t_size.wrapping_add(t_size.wrapping_neg() & data_align_minus_1)
+    }
+
+    /// Given a pointer to `data`, returns a pointer to `Struct`.
+    ///
+    /// # Safety
+    ///
+    /// The data pointer must have been allocated by `Self` with the same `data_layout`.
+    #[inline]
+    unsafe fn data_ptr_to_struct_ptr(data: NonNull<u8>, data_layout: Layout) -> NonNull<u8> {
+        unsafe {
+            let offset_of_data = Self::offset_of_data(data_layout);
+            NonNull::new_unchecked(data.as_ptr().sub(offset_of_data))
+        }
+    }
+
+    /// Given a pointer to `Struct`, returns a pointer to `data`.
+    ///
+    /// # Safety
+    ///
+    /// The struct pointer must have been allocated by `A` with the layout
+    /// `Self::struct_layout(data_layout)`.
+    #[inline]
+    unsafe fn struct_ptr_to_data_ptr(
+        struct_ptr: NonNull<[u8]>,
+        data_layout: Layout,
+    ) -> NonNull<[u8]> {
+        let offset_of_data = Self::offset_of_data(data_layout);
+        let data_ptr =
+            unsafe { NonNull::new_unchecked(struct_ptr.as_mut_ptr().add(offset_of_data)) };
+        // Note that the size is the exact size requested in the layout. Let me explain
+        // why this is necessary:
+        //
+        // Assume the original requested layout was `size=1, align=1`. Assume `T=u16`
+        // Then the struct layout is `size=4, align=2`. Assume that the allocator returns
+        // a slice with `size=5`. Then the space available for `data` is `3`.
+        // However, if we returned a slice with `len=3`, then the user would be allowed
+        // to call `dealloc` with `size=3, align=1`. In this case the struct layout
+        // would be computed as `size=6, align=2`. This size would be larger than what
+        // the allocator returned.
+        NonNull::slice_from_raw_parts(data_ptr, data_layout.size())
+    }
+}
+
+impl<T, A: Debug> Debug for StructAlloc<T, A> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        f.debug_struct("StructAlloc").field("0", &self.0).finish()
+    }
+}
+
+/// Delegates `Self::allocate{,_zereod}` to the allocator after computing the struct
+/// layout. Then transforms the new struct pointer to the new data pointer and returns it.
+macro delegate_alloc($id:ident) {
+    fn $id(&self, data_layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        let struct_layout = Self::struct_layout(data_layout)?;
+        let struct_ptr = self.0.$id(struct_layout)?;
+        unsafe { Ok(Self::struct_ptr_to_data_ptr(struct_ptr, data_layout)) }
+    }
+}
+
+/// Delegates `Self::{{grow{,_zeroed},shrink}` to the allocator after computing the struct
+/// layout and transforming the data pointer to the struct pointer. Then transforms
+/// the new struct pointer to the new data pointer and returns it.
+macro delegate_transform($id:ident) {
+    unsafe fn $id(
+        &self,
+        old_data_ptr: NonNull<u8>,
+        old_data_layout: Layout,
+        new_data_layout: Layout,
+    ) -> Result<NonNull<[u8]>, AllocError> {
+        let old_struct_layout = Self::struct_layout(old_data_layout)?;
+        let new_struct_layout = Self::struct_layout(new_data_layout)?;
+        unsafe {
+            let old_struct_ptr = Self::data_ptr_to_struct_ptr(old_data_ptr, old_data_layout);
+            let new_struct_ptr =
+                self.0.$id(old_struct_ptr, old_struct_layout, new_struct_layout)?;
+            Ok(Self::struct_ptr_to_data_ptr(new_struct_ptr, new_data_layout))
+        }
+    }
+}
+
+unsafe impl<T, A: Allocator> Allocator for StructAlloc<T, A> {
+    delegate_alloc!(allocate);
+    delegate_alloc!(allocate_zeroed);
+
+    unsafe fn deallocate(&self, data_ptr: NonNull<u8>, data_layout: Layout) {
+        unsafe {
+            let struct_ptr = Self::data_ptr_to_struct_ptr(data_ptr, data_layout);
+            let struct_layout =
+                Self::struct_layout(data_layout).expect("deallocate called with invalid layout");
+            self.0.deallocate(struct_ptr, struct_layout);
+        }
+    }
+
+    delegate_transform!(grow);
+    delegate_transform!(grow_zeroed);
+    delegate_transform!(shrink);
+}
+
+#[allow(unused_macros)]
+macro_rules! implement_struct_allocator {
+    ($id:ident) => {
+        #[unstable(feature = "struct_alloc", issue = "none")]
+        unsafe impl<A: Allocator> Allocator for $id<A> {
+            fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+                self.0.allocate(layout)
+            }
+
+            fn allocate_zeroed(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+                self.0.allocate_zeroed(layout)
+            }
+
+            unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+                unsafe { self.0.deallocate(ptr, layout) }
+            }
+
+            unsafe fn grow(
+                &self,
+                ptr: NonNull<u8>,
+                old_layout: Layout,
+                new_layout: Layout,
+            ) -> Result<NonNull<[u8]>, AllocError> {
+                unsafe { self.0.grow(ptr, old_layout, new_layout) }
+            }
+
+            unsafe fn grow_zeroed(
+                &self,
+                ptr: NonNull<u8>,
+                old_layout: Layout,
+                new_layout: Layout,
+            ) -> Result<NonNull<[u8]>, AllocError> {
+                unsafe { self.0.grow_zeroed(ptr, old_layout, new_layout) }
+            }
+
+            unsafe fn shrink(
+                &self,
+                ptr: NonNull<u8>,
+                old_layout: Layout,
+                new_layout: Layout,
+            ) -> Result<NonNull<[u8]>, AllocError> {
+                unsafe { self.0.shrink(ptr, old_layout, new_layout) }
+            }
+        }
+    };
+}

library/alloc/src/alloc/struct_alloc/tests.rs

@@ -0,0 +1,126 @@
+use super::StructAlloc;
+use crate::alloc::Global;
+use crate::collections::VecDeque;
+use core::alloc::{AllocError, Allocator};
+use std::alloc::Layout;
+use std::cell::RefCell;
+use std::ptr::NonNull;
+use std::{any, ptr};
+
+fn test_pair<T, Data>() {
+    if let Err(_) = std::panic::catch_unwind(|| test_pair_::<T, Data>()) {
+        panic!("test of {} followed by {} failed", any::type_name::<T>(), any::type_name::<Data>());
+    }
+}
+
+fn test_pair_<T, Data>() {
+    #[repr(C)]
+    struct S<T, Data> {
+        t: T,
+        data: Data,
+    }
+
+    let offset = {
+        let s: *const S<T, Data> = ptr::null();
+        unsafe { std::ptr::addr_of!((*s).data) as usize }
+    };
+
+    let expected_layout = RefCell::new(VecDeque::new());
+    let expected_ptr = RefCell::new(VecDeque::new());
+
+    let check_layout = |actual| {
+        let mut e = expected_layout.borrow_mut();
+        match e.pop_front() {
+            Some(expected) if expected == actual => {}
+            Some(expected) => panic!("expected layout {:?}, actual layout {:?}", expected, actual),
+            _ => panic!("unexpected allocator invocation with layout {:?}", actual),
+        }
+    };
+
+    let check_ptr = |actual: NonNull<u8>| {
+        let mut e = expected_ptr.borrow_mut();
+        match e.pop_front() {
+            Some(expected) if expected == actual.as_ptr() => {}
+            Some(expected) => {
+                panic!("expected pointer {:p}, actual pointer {:p}", expected, actual)
+            }
+            _ => panic!("unexpected allocator invocation with pointer {:p}", actual),
+        }
+    };
+
+    struct TestAlloc<F, G>(F, G);
+
+    unsafe impl<F, G> Allocator for TestAlloc<F, G>
+    where
+        F: Fn(Layout),
+        G: Fn(NonNull<u8>),
+    {
+        fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+            self.0(layout);
+            Global.allocate(layout)
+        }
+
+        unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+            self.1(ptr);
+            self.0(layout);
+            unsafe { Global.deallocate(ptr, layout) }
+        }
+    }
+
+    let struct_alloc = StructAlloc::<T, _>::new(TestAlloc(check_layout, check_ptr));
+
+    fn s_layout<T, Data, const N: usize>() -> Layout {
+        Layout::new::<S<T, [Data; N]>>()
+    }
+
+    fn d_layout<Data, const N: usize>() -> Layout {
+        Layout::new::<[Data; N]>()
+    }
+
+    fn check_slice<Data, const N: usize>(ptr: NonNull<[u8]>) {
+        let expected = d_layout::<Data, N>().size();
+        if ptr.len() != expected {
+            panic!(
+                "expected allocation size: {:?}, actual allocation size: {:?}",
+                expected,
+                ptr.len()
+            )
+        }
+    }
+
+    expected_layout.borrow_mut().push_back(s_layout::<T, Data, 1>());
+    let ptr = struct_alloc.allocate(d_layout::<Data, 1>()).unwrap();
+    check_slice::<Data, 1>(ptr);
+    unsafe {
+        expected_ptr.borrow_mut().push_back(ptr.as_mut_ptr().sub(offset));
+    }
+    expected_layout.borrow_mut().push_back(s_layout::<T, Data, 3>());
+    expected_layout.borrow_mut().push_back(s_layout::<T, Data, 1>());
+    let ptr = unsafe {
+        struct_alloc
+            .grow(ptr.as_non_null_ptr(), d_layout::<Data, 1>(), d_layout::<Data, 3>())
+            .unwrap()
+    };
+    check_slice::<Data, 3>(ptr);
+    unsafe {
+        expected_ptr.borrow_mut().push_back(ptr.as_mut_ptr().sub(offset));
+    }
+    expected_layout.borrow_mut().push_back(s_layout::<T, Data, 3>());
+    unsafe {
+        struct_alloc.deallocate(ptr.as_non_null_ptr(), d_layout::<Data, 3>());
+    }
+    if !expected_ptr.borrow().is_empty() || !expected_layout.borrow().is_empty() {
+        panic!("missing allocator calls");
+    }
+}
+
+#[test]
+fn test() {
+    macro_rules! test_ty {
+        ($($ty:ty),*) => { test_ty!(@2 $($ty),*; ($($ty),*)) };
+        (@2 $($tyl:ty),*; $tyr:tt) => { $(test_ty!(@3 $tyl; $tyr);)* };
+        (@3 $tyl:ty; ($($tyr:ty),*)) => { $(test_pair::<$tyl, $tyr>();)* };
+    }
+    // call test_pair::<A, B>() for every combination of these types
+    test_ty!((), u8, u16, u32, u64, u128);
+}

library/alloc/src/lib.rs

@@ -155,6 +155,7 @@ mod macros;
 
 // Heaps provided for low-level allocation strategies
 
+#[macro_use]
 pub mod alloc;
 
 // Primitive types using the heaps above