Add Allocator-helper to allocate prefix and safety-tester

Author: TimDiekmann

library/core/src/alloc/helper.rs

@@ -0,0 +1,170 @@
+use crate::{
+    alloc::{AllocError, Allocator, Layout},
+    fmt,
+    marker::PhantomData,
+    ptr::NonNull,
+};
+
+/// An allocator that requests some extra memory from the parent allocator for storing a prefix and/or a suffix.
+///
+/// The alignment of the memory block is the maximum of the alignment of `Prefix` and the requested
+/// alignment. This may introduce an unused padding between `Prefix` and the returned memory.
+///
+/// To get a pointer to the prefix, [`prefix()`] may be called.
+///
+/// [`prefix()`]: Self::prefix
+///
+/// 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>()`.
+#[unstable(feature = "allocator_api_internals", issue = "none")]
+pub struct PrefixAllocator<Alloc, Prefix = ()> {
+    /// The parent allocator to be used as backend
+    pub parent: Alloc,
+    _prefix: PhantomData<*const Prefix>,
+}
+
+impl<Alloc: fmt::Debug, Prefix> fmt::Debug for PrefixAllocator<Alloc, Prefix> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("Affix").field("parent", &self.parent).finish()
+    }
+}
+
+impl<Alloc: Default, Prefix> Default for PrefixAllocator<Alloc, Prefix> {
+    fn default() -> Self {
+        Self::new(Alloc::default())
+    }
+}
+
+impl<Alloc: Clone, Prefix> Clone for PrefixAllocator<Alloc, Prefix> {
+    fn clone(&self) -> Self {
+        Self::new(self.parent.clone())
+    }
+}
+
+impl<Alloc: Copy, Prefix> Copy for PrefixAllocator<Alloc, Prefix> {}
+
+impl<Alloc: PartialEq, Prefix> PartialEq for PrefixAllocator<Alloc, Prefix> {
+    fn eq(&self, other: &Self) -> bool {
+        self.parent.eq(&other.parent)
+    }
+}
+
+impl<Alloc: Eq, Prefix> Eq for PrefixAllocator<Alloc, Prefix> {}
+
+unsafe impl<Alloc: Send, Prefix> Send for PrefixAllocator<Alloc, Prefix> {}
+unsafe impl<Alloc: Sync, Prefix> Sync for PrefixAllocator<Alloc, Prefix> {}
+impl<Alloc: Unpin, Prefix> Unpin for PrefixAllocator<Alloc, Prefix> {}
+
+impl<Alloc, Prefix> PrefixAllocator<Alloc, Prefix> {
+    pub const fn new(parent: Alloc) -> Self {
+        Self { parent, _prefix: PhantomData }
+    }
+
+    /// Returns the offset between the `Prefix` and the stored data.
+    #[inline]
+    pub fn prefix_offset(layout: Layout) -> usize {
+        let prefix_layout = Layout::new::<Prefix>();
+        prefix_layout.size() + prefix_layout.padding_needed_for(layout.align())
+    }
+
+    /// Returns a pointer to the prefix.
+    ///
+    /// # Safety
+    ///
+    /// * `ptr` must denote a block of memory *[currently allocated]* via this allocator, and
+    /// * `layout` must *[fit]* that block of memory.
+    ///
+    /// [currently allocated]: https://doc.rust-lang.org/nightly/core/alloc/trait.AllocRef.html#currently-allocated-memory
+    /// [fit]: https://doc.rust-lang.org/nightly/core/alloc/trait.AllocRef.html#memory-fitting
+    #[inline]
+    pub unsafe fn prefix(ptr: NonNull<u8>, layout: Layout) -> NonNull<Prefix> {
+        let prefix_offset = Self::prefix_offset(layout);
+        // SAFETY: `prefix_offset` is smaller (and not equal to) `ptr` as the same function for calculating `prefix_offset` is used when allocating.
+        unsafe { NonNull::new_unchecked(ptr.as_ptr().sub(prefix_offset)).cast() }
+    }
+
+    fn create_ptr(ptr: NonNull<[u8]>, offset_prefix: usize) -> NonNull<[u8]> {
+        let len = ptr.len() - offset_prefix;
+
+        // SAFETY: `prefix_offset` is smaller (and not equal to) `ptr` as the same function for calculating `prefix_offset` is used when allocating.
+        let ptr = unsafe { NonNull::new_unchecked(ptr.as_mut_ptr().add(offset_prefix)) };
+
+        NonNull::slice_from_raw_parts(ptr, len)
+    }
+
+    #[inline]
+    fn alloc_impl(
+        layout: Layout,
+        alloc: impl FnOnce(Layout) -> Result<NonNull<[u8]>, AllocError>,
+    ) -> Result<NonNull<[u8]>, AllocError> {
+        let (layout, offset_prefix) =
+            Layout::new::<Prefix>().extend(layout).map_err(|_| AllocError)?;
+
+        Ok(Self::create_ptr(alloc(layout)?, offset_prefix))
+    }
+}
+
+unsafe impl<Alloc, Prefix> Allocator for PrefixAllocator<Alloc, Prefix>
+where
+    Alloc: Allocator,
+{
+    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        Self::alloc_impl(layout, |l| self.parent.allocate(l))
+    }
+
+    fn allocate_zeroed(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        Self::alloc_impl(layout, |l| self.parent.allocate_zeroed(l))
+    }
+
+    unsafe fn grow(
+        &self,
+        _ptr: NonNull<u8>,
+        _old_layout: Layout,
+        _new_layout: Layout,
+    ) -> Result<NonNull<[u8]>, AllocError> {
+        // For (A)Rc it's not needed. When implementing please take care, if the alignment changes.
+        unimplemented!("PrefixAllocator currently does not implement growing.");
+    }
+
+    unsafe fn grow_zeroed(
+        &self,
+        _ptr: NonNull<u8>,
+        _old_layout: Layout,
+        _new_layout: Layout,
+    ) -> Result<NonNull<[u8]>, AllocError> {
+        // For (A)Rc it's not needed. When implementing please take care, if the alignment changes.
+        unimplemented!("PrefixAllocator currently does not implement growing.");
+    }
+
+    unsafe fn shrink(
+        &self,
+        _ptr: NonNull<u8>,
+        _old_layout: Layout,
+        _new_layout: Layout,
+    ) -> Result<NonNull<[u8]>, AllocError> {
+        // For (A)Rc it's not needed. When implementing please take care, if the alignment changes.
+        unimplemented!("PrefixAllocator currently does not implement shrinking.");
+    }
+
+    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+        let (layout, prefix_offset) = Layout::new::<Prefix>().extend(layout).unwrap();
+        // SAFETY: `prefix_offset` is smaller (and not equal to) `ptr` as the same function for calculating `prefix_offset` is used when allocating.
+        unsafe {
+            let base_ptr = NonNull::new_unchecked(ptr.as_ptr().sub(prefix_offset));
+            self.parent.deallocate(base_ptr, layout)
+        };
+    }
+}

library/core/src/alloc/mod.rs

@@ -5,6 +5,9 @@
 mod global;
 mod layout;
 
+#[unstable(feature = "allocator_api_internals", issue = "none")]
+pub mod helper;
+
 #[stable(feature = "global_alloc", since = "1.28.0")]
 pub use self::global::GlobalAlloc;
 #[stable(feature = "alloc_layout", since = "1.28.0")]

library/core/tests/alloc.rs renamed to library/core/tests/alloc/const_unchecked_layout.rs

@@ -0,0 +1,185 @@
+use core::alloc::{AllocError, Allocator, Layout};
+use core::ptr::NonNull;
+use std::{
+    collections::HashMap,
+    sync::{Mutex, PoisonError},
+};
+
+mod const_unchecked_layout;
+mod prefix;
+
+#[derive(Default)]
+/// Implements `Allocator` and checks it's unsafety conditions.
+struct Tracker<A> {
+    alloc: A,
+    map: Mutex<HashMap<NonNull<u8>, (usize, Layout)>>,
+}
+
+impl<A> Tracker<A> {
+    fn new(alloc: A) -> Self {
+        Self { alloc, map: Default::default() }
+    }
+
+    fn after_alloc(&self, layout: Layout, result: Result<NonNull<[u8]>, AllocError>) {
+        if let Ok(ptr) = result {
+            self.map
+                .lock()
+                .unwrap_or_else(PoisonError::into_inner)
+                .insert(ptr.as_non_null_ptr(), (ptr.len(), layout));
+        }
+    }
+
+    fn before_dealloc(&self, ptr: NonNull<u8>, layout: Layout) {
+        let lock = self.map.lock().unwrap_or_else(PoisonError::into_inner);
+        let (size, old_layout) = lock
+            .get(&ptr)
+            .expect("`ptr` must denote a block of memory currently allocated via this allocator");
+        assert_eq!(
+            layout.align(),
+            old_layout.align(),
+            "`layout` must fit that block of memory. Expected alignment of {}, got {}",
+            old_layout.align(),
+            layout.align()
+        );
+        if layout.size() < old_layout.size() || layout.size() > *size {
+            if *size == old_layout.size() {
+                panic!(
+                    "`layout` must fit that block of memory. Expected size of {}, got {}",
+                    old_layout.size(),
+                    layout.size()
+                )
+            } else {
+                panic!(
+                    "`layout` must fit that block of memory. Expected size between {}..={}, \
+                        got {}",
+                    old_layout.size(),
+                    size,
+                    layout.size()
+                )
+            }
+        }
+    }
+
+    fn after_dealloc(&self, ptr: NonNull<u8>, _layout: Layout) {
+        self.map.lock().unwrap_or_else(PoisonError::into_inner).remove(&ptr);
+    }
+
+    fn before_grow(&self, ptr: NonNull<u8>, old_layout: Layout, new_layout: Layout) {
+        assert!(
+            new_layout.size() >= old_layout.size(),
+            "`new_layout.size()` must be greater than or equal to `old_layout.size()`, expected {} >= {}",
+            new_layout.size(),
+            old_layout.size()
+        );
+        self.before_dealloc(ptr, old_layout)
+    }
+
+    #[track_caller]
+    fn after_grow(
+        &self,
+        ptr: NonNull<u8>,
+        old_layout: Layout,
+        new_layout: Layout,
+        result: Result<NonNull<[u8]>, AllocError>,
+    ) {
+        if result.is_ok() {
+            self.after_dealloc(ptr, old_layout);
+            self.after_alloc(new_layout, result);
+        }
+    }
+
+    fn before_shrink(&self, ptr: NonNull<u8>, old_layout: Layout, new_layout: Layout) {
+        assert!(
+            new_layout.size() <= old_layout.size(),
+            "`new_layout.size()` must be smaller than or equal to `old_layout.size()`, expected {} >= {}",
+            new_layout.size(),
+            old_layout.size()
+        );
+        self.before_dealloc(ptr, old_layout)
+    }
+
+    #[track_caller]
+    fn after_shrink(
+        &self,
+        ptr: NonNull<u8>,
+        old_layout: Layout,
+        new_layout: Layout,
+        result: Result<NonNull<[u8]>, AllocError>,
+    ) {
+        if result.is_ok() {
+            self.after_dealloc(ptr, old_layout);
+            self.after_alloc(new_layout, result);
+        }
+    }
+}
+
+unsafe impl<A: Allocator> Allocator for Tracker<A> {
+    #[track_caller]
+    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, std::alloc::AllocError> {
+        let result = self.alloc.allocate(layout);
+        self.after_alloc(layout, result);
+        result
+    }
+
+    #[track_caller]
+    fn allocate_zeroed(&self, layout: Layout) -> Result<NonNull<[u8]>, std::alloc::AllocError> {
+        let result = self.alloc.allocate_zeroed(layout);
+        self.after_alloc(layout, result);
+        result
+    }
+
+    #[track_caller]
+    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+        self.before_dealloc(ptr, layout);
+        unsafe { self.alloc.deallocate(ptr, layout) }
+        self.after_dealloc(ptr, layout);
+    }
+
+    #[track_caller]
+    unsafe fn grow(
+        &self,
+        ptr: NonNull<u8>,
+        old_layout: Layout,
+        new_layout: Layout,
+    ) -> Result<NonNull<[u8]>, std::alloc::AllocError> {
+        self.before_grow(ptr, old_layout, new_layout);
+        let result = unsafe { self.alloc.grow(ptr, old_layout, new_layout) };
+        self.after_grow(ptr, old_layout, new_layout, result);
+        result
+    }
+
+    #[track_caller]
+    unsafe fn grow_zeroed(
+        &self,
+        ptr: NonNull<u8>,
+        old_layout: Layout,
+        new_layout: Layout,
+    ) -> Result<NonNull<[u8]>, std::alloc::AllocError> {
+        self.before_grow(ptr, old_layout, new_layout);
+        let result = unsafe { self.alloc.grow_zeroed(ptr, old_layout, new_layout) };
+        self.after_grow(ptr, old_layout, new_layout, result);
+        result
+    }
+
+    #[track_caller]
+    unsafe fn shrink(
+        &self,
+        ptr: NonNull<u8>,
+        old_layout: Layout,
+        new_layout: Layout,
+    ) -> Result<NonNull<[u8]>, std::alloc::AllocError> {
+        self.before_shrink(ptr, old_layout, new_layout);
+        let result = unsafe { self.alloc.shrink(ptr, old_layout, new_layout) };
+        self.after_shrink(ptr, old_layout, new_layout, result);
+        result
+    }
+}
+
+impl<A> Drop for Tracker<A> {
+    fn drop(&mut self) {
+        let lock = self.map.lock().unwrap_or_else(PoisonError::into_inner);
+        if !lock.is_empty() {
+            panic!("Missing deallocations {:#?}", lock);
+        }
+    }
+}