Restructure Rc and Arc metadata for uniform access via a PrefixAlloc

Author: TimDiekmann

library/alloc/src/lib.rs

@@ -76,6 +76,7 @@
 #![cfg_attr(test, feature(test))]
 #![cfg_attr(test, feature(new_uninit))]
 #![feature(allocator_api)]
+#![feature(allocator_api_internals)]
 #![feature(vec_extend_from_within)]
 #![feature(array_chunks)]
 #![feature(array_methods)]

library/alloc/src/rc.rs

@@ -247,6 +247,7 @@ use crate::boxed::Box;
 #[cfg(test)]
 use std::boxed::Box;
 
+use core::alloc::helper::PrefixAllocator;
 use core::any::Any;
 use core::borrow;
 use core::cell::Cell;
@@ -257,7 +258,7 @@ use core::hash::{Hash, Hasher};
 use core::intrinsics::abort;
 use core::iter;
 use core::marker::{self, PhantomData, Unpin, Unsize};
-use core::mem::{self, align_of_val_raw, forget, size_of_val};
+use core::mem::{self, forget, size_of_val, MaybeUninit};
 use core::ops::{CoerceUnsized, Deref, DispatchFromDyn, Receiver};
 use core::pin::Pin;
 use core::ptr::{self, NonNull};
@@ -273,13 +274,33 @@ use crate::vec::Vec;
 #[cfg(test)]
 mod tests;
 
+struct RcBoxMetadata {
+    strong: Cell<usize>,
+    weak: Cell<usize>,
+}
+
+impl RcBoxMetadata {
+    // There is an implicit weak pointer owned by all the strong
+    // pointers, which ensures that the weak destructor never frees
+    // the allocation while the strong destructor is running, even
+    // if the weak pointer is stored inside the strong one.
+    #[inline]
+    fn new_strong() -> Self {
+        Self { strong: Cell::new(1), weak: Cell::new(1) }
+    }
+
+    #[inline]
+    fn new_weak() -> Self {
+        Self { strong: Cell::new(0), weak: Cell::new(1) }
+    }
+}
+
 // This is repr(C) to future-proof against possible field-reordering, which
 // would interfere with otherwise safe [into|from]_raw() of transmutable
 // inner types.
 #[repr(C)]
 struct RcBox<T: ?Sized> {
-    strong: Cell<usize>,
-    weak: Cell<usize>,
+    meta: RcBoxMetadata,
     value: T,
 }
 
@@ -319,10 +340,12 @@ impl<T: ?Sized> Rc<T> {
         unsafe { self.ptr.as_ref() }
     }
 
+    #[inline]
     fn from_inner(ptr: NonNull<RcBox<T>>) -> Self {
         Self { ptr, phantom: PhantomData }
     }
 
+    #[inline]
     unsafe fn from_ptr(ptr: *mut RcBox<T>) -> Self {
         Self::from_inner(unsafe { NonNull::new_unchecked(ptr) })
     }
@@ -340,13 +363,7 @@ impl<T> Rc<T> {
     /// ```
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn new(value: T) -> Rc<T> {
-        // There is an implicit weak pointer owned by all the strong
-        // pointers, which ensures that the weak destructor never frees
-        // the allocation while the strong destructor is running, even
-        // if the weak pointer is stored inside the strong one.
-        Self::from_inner(
-            Box::leak(box RcBox { strong: Cell::new(1), weak: Cell::new(1), value }).into(),
-        )
+        Self::from_inner(Box::leak(box RcBox { meta: RcBoxMetadata::new_strong(), value }).into())
     }
 
     /// Constructs a new `Rc<T>` using a weak reference to itself. Attempting
@@ -378,8 +395,7 @@ impl<T> Rc<T> {
         // Construct the inner in the "uninitialized" state with a single
         // weak reference.
         let uninit_ptr: NonNull<_> = Box::leak(box RcBox {
-            strong: Cell::new(0),
-            weak: Cell::new(1),
+            meta: RcBoxMetadata::new_weak(),
             value: mem::MaybeUninit::<T>::uninit(),
         })
         .into();
@@ -400,9 +416,9 @@ impl<T> Rc<T> {
             let inner = init_ptr.as_ptr();
             ptr::write(ptr::addr_of_mut!((*inner).value), data);
 
-            let prev_value = (*inner).strong.get();
+            let prev_value = (*inner).meta.strong.get();
             debug_assert_eq!(prev_value, 0, "No prior strong references should exist");
-            (*inner).strong.set(1);
+            (*inner).meta.strong.set(1);
         }
 
         let strong = Rc::from_inner(init_ptr);
@@ -494,8 +510,7 @@ impl<T> Rc<T> {
         // the allocation while the strong destructor is running, even
         // if the weak pointer is stored inside the strong one.
         Ok(Self::from_inner(
-            Box::leak(Box::try_new(RcBox { strong: Cell::new(1), weak: Cell::new(1), value })?)
-                .into(),
+            Box::leak(Box::try_new(RcBox { meta: RcBoxMetadata::new_strong(), value })?).into(),
         ))
     }
 
@@ -846,13 +861,7 @@ impl<T: ?Sized> Rc<T> {
     /// ```
     #[stable(feature = "rc_raw", since = "1.17.0")]
     pub unsafe fn from_raw(ptr: *const T) -> Self {
-        let offset = unsafe { data_offset(ptr) };
-
-        // Reverse the offset to find the original RcBox.
-        let rc_ptr =
-            unsafe { (ptr as *mut RcBox<T>).set_ptr_value((ptr as *mut u8).offset(-offset)) };
-
-        unsafe { Self::from_ptr(rc_ptr) }
+        unsafe { Self::from_data_ptr(ptr).assume_init() }
     }
 
     /// Creates a new [`Weak`] pointer to this allocation.
@@ -1237,8 +1246,8 @@ impl<T: ?Sized> Rc<T> {
         unsafe {
             debug_assert_eq!(Layout::for_value(&*inner), layout);
 
-            ptr::write(&mut (*inner).strong, Cell::new(1));
-            ptr::write(&mut (*inner).weak, Cell::new(1));
+            ptr::write(&mut (*inner).meta.strong, Cell::new(1));
+            ptr::write(&mut (*inner).meta.weak, Cell::new(1));
         }
 
         Ok(inner)
@@ -1277,6 +1286,23 @@ impl<T: ?Sized> Rc<T> {
             Self::from_ptr(ptr)
         }
     }
+
+    /// # Safety
+    ///
+    /// The caller must ensure that the pointer points to the `value` field of a `Global`
+    /// allocation of type `RcBox<T>`. Depending on how the pointer was created, the
+    /// `meta` field might or might not be uninitialized. It's up to the caller to ensure
+    /// that this field is set to the correct value before the return value is unwrapped.
+    #[inline]
+    unsafe fn from_data_ptr(ptr: *const T) -> MaybeUninit<Self> {
+        let offset = unsafe { data_offset(ptr) };
+
+        // Reverse the offset to find the original RcBox.
+        let rc_ptr =
+            unsafe { (ptr as *mut RcBox<T>).set_ptr_value((ptr as *mut u8).offset(-offset)) };
+
+        unsafe { MaybeUninit::new(Self::from_ptr(rc_ptr)) }
+    }
 }
 
 impl<T> Rc<[T]> {
@@ -2206,7 +2232,7 @@ impl<T: ?Sized> Weak<T> {
             // is dropped, the data field will be dropped in-place).
             Some(unsafe {
                 let ptr = self.ptr.as_ptr();
-                WeakInner { strong: &(*ptr).strong, weak: &(*ptr).weak }
+                WeakInner { strong: &(*ptr).meta.strong, weak: &(*ptr).meta.weak }
             })
         }
     }
@@ -2415,12 +2441,12 @@ trait RcInnerPtr {
 impl<T: ?Sized> RcInnerPtr for RcBox<T> {
     #[inline(always)]
     fn weak_ref(&self) -> &Cell<usize> {
-        &self.weak
+        &self.meta.weak
     }
 
     #[inline(always)]
     fn strong_ref(&self) -> &Cell<usize> {
-        &self.strong
+        &self.meta.strong
     }
 }
 
@@ -2453,24 +2479,21 @@ impl<T: ?Sized> AsRef<T> for Rc<T> {
 #[stable(feature = "pin", since = "1.33.0")]
 impl<T: ?Sized> Unpin for Rc<T> {}
 
+type RcAllocator = PrefixAllocator<RcBoxMetadata, Global>;
+
 /// Get the offset within an `RcBox` for the payload behind a pointer.
 ///
 /// # Safety
 ///
 /// The pointer must point to (and have valid metadata for) a previously
 /// valid instance of T, but the T is allowed to be dropped.
 unsafe fn data_offset<T: ?Sized>(ptr: *const T) -> isize {
-    // Align the unsized value to the end of the RcBox.
-    // Because RcBox is repr(C), it will always be the last field in memory.
-    // SAFETY: since the only unsized types possible are slices, trait objects,
-    // and extern types, the input safety requirement is currently enough to
-    // satisfy the requirements of align_of_val_raw; this is an implementation
-    // detail of the language that may not be relied upon outside of std.
-    unsafe { data_offset_align(align_of_val_raw(ptr)) }
-}
-
-#[inline]
-fn data_offset_align(align: usize) -> isize {
-    let layout = Layout::new::<RcBox<()>>();
-    (layout.size() + layout.padding_needed_for(align)) as isize
+    unsafe {
+        // SAFETY: since the only unsized types possible are slices, trait objects,
+        // and extern types, the input safety requirement is currently enough to
+        // satisfy the requirements of for_value_raw; this is an implementation
+        // detail of the language that may not be relied upon outside of std.
+        let layout = Layout::for_value_raw(ptr);
+        RcAllocator::prefix_offset(layout) as isize
+    }
 }