Consider the allocator stored in an allocation

Avoiding the pointer to the allocator in every structure that contains a
non-global allocator is possible by storing the pointer inside scratch
space of existing allocations. However, that requires being passed one.

Author: 197g

alloc-traits/src/ref_.rs

@@ -4,65 +4,97 @@ use crate::NonZeroLayout;
 use crate::LocalAlloc;
 
 /// An allocation without tracked lifetime.
+#[derive(Copy, Clone)]
 pub struct Allocation {
     ptr: NonNull<u8>,
     layout: NonZeroLayout,
 }
 
 /// A references to an allocator.
 ///
-/// Allocations must be live for the lifetime of `self`. That is one must be able to move the
-/// `self` and keep all allocations. In particular, a reference to a [`LocalAlloc`] is an
-/// `AllocRef`.
+/// Allocations created from this instance are valid for the lifetime of `self`. That is one must
+/// be able to move the `self` and all allocations remain valid. In particular, a reference to a
+/// [`LocalAlloc`] is an `AllocRef` as it is an immutable pin.
+///
+/// An allocation is said to be live when it is valid and has not been passed to `dealloc` and has
+/// not been passed to `realloc` that returned `Some(_)`.
+///
+/// Note that all methods require being passed some existing allocation, even if they do not
+/// consume it. See `InitialAllocRef` for those that do not.
 ///
 /// [`LocalAlloc`]: trait.LocalAlloc.html
 pub unsafe trait AllocRef {
     /// Allocate one block of memory.
     ///
-    /// The callee guarantees that a successful return contains a pointer that is valid for **at
-    /// least** the layout requested by the caller.
-    fn alloc(&mut self, layout: NonZeroLayout) -> Option<Allocation>;
-
-    /// Deallocate a block previously allocated.
-    /// # Safety
-    /// The caller must ensure that:
-    /// * `alloc` has been previously returned from a call to `alloc`.
-    /// * There are no more pointer to the allocation.
-    unsafe fn dealloc(&mut self, alloc: Allocation);
+    /// The `existing` allocation can be used by the `AllocRef` to uniquely identify the allocator
+    /// that is used internaly. It must be valid for the same `AllocRef`.
+    ///
+    /// The callee guarantees that a successful return contains a pointer that fits the layout
+    /// requested by the caller and the layout in the struct is the same as requested.
+    unsafe fn alloc_from(&mut self, existing: Allocation, layout: NonZeroLayout) -> Option<Allocation>;
 
     /// Allocate a block of memory initialized with zeros.
     ///
-    /// The callee guarantees that a successful return contains a pointer that is valid for **at
-    /// least** the layout requested by the caller and the contiguous region of bytes, starting at
-    /// the pointer and with the size of the returned layout, is initialized and zeroed.
-    fn alloc_zeroed(&mut self, layout: NonZeroLayout)
+    /// The callee guarantees that a successful return contains a pointer that fits the layout
+    /// requested by the caller and the layout in the struct is the same as requested.
+    unsafe fn alloc_zeroed_from(&mut self, existing: Allocation, layout: NonZeroLayout)
         -> Option<Allocation> 
     {
-        let allocation = self.alloc(layout)?;
-        unsafe {
-            write_bytes(allocation.ptr.as_ptr(), 0u8, allocation.layout.size().into());
-        }
+        let allocation = self.alloc_from(existing, layout)?;
+        write_bytes(allocation.ptr.as_ptr(), 0u8, allocation.layout.size().into());
         Some(allocation)
     }
 
+    /// Deallocate a block previously allocated.
+    /// # Safety
+    /// The caller must ensure that:
+    /// * `alloc` has been previously returned from a call to `alloc`.
+    /// * There are no more pointer to the allocation.
+    unsafe fn dealloc(&mut self, alloc: Allocation);
+
     /// Change the layout of a block previously allocated.
     ///
-    /// The callee guarantees that a successful return contains a pointer that is valid for **at
-    /// least** the layout requested by the caller and the contiguous region of bytes, starting at
-    /// the pointer and with the size of the returned layout, is initialized with the prefix of the
-    /// previous allocation that is still valid.
+    /// The callee guarantees that a successful return contains a pointer that it fits the layout
+    /// requested by the caller and the contiguous region of bytes, starting at the pointer and
+    /// with the size of the returned layout, is initialized with the prefix of the previous
+    /// allocation that is still valid, and that the layout in the returned struct is the same as
+    /// requested.
     unsafe fn realloc(&mut self, alloc: Allocation, layout: NonZeroLayout)
         -> Option<Allocation>
     {
-        let new_alloc = self.alloc(layout)?;
+        let new_alloc = self.alloc_from(alloc, layout)?;
         copy_nonoverlapping(
             alloc.ptr.as_ptr(),
             new_alloc.ptr.as_ptr(),
             layout.size().min(alloc.layout.size()).into());
+        self.dealloc(alloc);
         Some(new_alloc)
     }
 }
 
+/// A trait for an `AllocRef` that requires no existing allocation to allocate.
+pub unsafe trait InitialAllocRef: AllocRef {
+    /// Allocate one block of memory.
+    ///
+    /// The callee guarantees that a successful return contains a pointer that fits the layout
+    /// requested by the caller and the layout in the struct is the same as requested.
+    fn alloc(&mut self, layout: NonZeroLayout) -> Option<Allocation>;
+
+    /// Allocate a block of memory initialized with zeros.
+    ///
+    /// The callee guarantees that a successful return contains a pointer that fits the layout
+    /// requested by the caller and the layout in the struct is the same as requested.
+    fn alloc_zeroed(&mut self, layout: NonZeroLayout)
+        -> Option<Allocation> 
+    {
+        let allocation = InitialAllocRef::alloc(self, layout)?;
+        unsafe {
+            write_bytes(allocation.ptr.as_ptr(), 0u8, allocation.layout.size().into());
+        }
+        Some(allocation)
+    }
+}
+
 impl Allocation {
     pub(crate) fn from_local(alloc: crate::local::Allocation) -> Self {
         Allocation {
@@ -89,18 +121,18 @@ impl Allocation {
 unsafe impl<'rf, T> AllocRef for &'rf T
     where T: LocalAlloc<'rf>
 {
-    fn alloc(&mut self, layout: NonZeroLayout) -> Option<Allocation> {
-        LocalAlloc::<'rf>::alloc(*self, layout).map(Allocation::from_local)
+    unsafe fn alloc_from(&mut self, _: Allocation, layout: NonZeroLayout) -> Option<Allocation> {
+        InitialAllocRef::alloc(self, layout)
     }
 
-    unsafe fn dealloc(&mut self, alloc: Allocation) {
-        LocalAlloc::<'rf>::dealloc(*self, alloc.into_local())
-    }
-
-    fn alloc_zeroed(&mut self, layout: NonZeroLayout)
+    unsafe fn alloc_zeroed_from(&mut self, _: Allocation, layout: NonZeroLayout)
         -> Option<Allocation> 
     {
-        LocalAlloc::<'rf>::alloc_zeroed(*self, layout).map(Allocation::from_local)
+        InitialAllocRef::alloc_zeroed(self, layout)
+    }
+
+    unsafe fn dealloc(&mut self, alloc: Allocation) {
+        LocalAlloc::<'rf>::dealloc(*self, alloc.into_local())
     }
 
     unsafe fn realloc(&mut self, alloc: Allocation, layout: NonZeroLayout)
@@ -110,3 +142,15 @@ unsafe impl<'rf, T> AllocRef for &'rf T
             .map(Allocation::from_local)
     }
 }
+
+unsafe impl<'rf, T> InitialAllocRef for &'rf T
+    where T: LocalAlloc<'rf>
+{
+    fn alloc(&mut self, layout: NonZeroLayout) -> Option<Allocation> {
+        LocalAlloc::<'rf>::alloc(*self, layout).map(Allocation::from_local)
+    }
+
+    fn alloc_zeroed(&mut self, layout: NonZeroLayout) -> Option<Allocation> {
+        LocalAlloc::<'rf>::alloc_zeroed(*self, layout).map(Allocation::from_local)
+    }
+}