rust-lang · matthiaskrgr · Feb 17, 2024 · Jun 11, 2024 · Jun 12, 2024 · Jun 12, 2024
diff --git a/library/alloc/src/rc.rs b/library/alloc/src/rc.rs
@@ -3516,7 +3516,7 @@ fn data_offset_align(align: usize) -> usize {
     layout.size() + layout.padding_needed_for(align)
 }
 
-/// A uniquely owned `Rc`
+/// A uniquely owned [`Rc`].
 ///
 /// This represents an `Rc` that is known to be uniquely owned -- that is, have exactly one strong
 /// reference. Multiple weak pointers can be created, but attempts to upgrade those to strong
@@ -3554,13 +3554,24 @@ fn data_offset_align(align: usize) -> usize {
 /// including fallible or async constructors.
 #[unstable(feature = "unique_rc_arc", issue = "112566")]
 #[derive(Debug)]
-pub struct UniqueRc<T> {
+pub struct UniqueRc<
+    T: ?Sized,
+    #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator = Global,
+> {
     ptr: NonNull<RcBox<T>>,
     phantom: PhantomData<RcBox<T>>,
+    alloc: A,
 }
 
+#[unstable(feature = "unique_rc_arc", issue = "112566")]
+impl<T: ?Sized + Unsize<U>, U: ?Sized, A: Allocator> CoerceUnsized<UniqueRc<U, A>>
+    for UniqueRc<T, A>
+{
+}
+
+// Depends on A = Global
 impl<T> UniqueRc<T> {
-    /// Creates a new `UniqueRc`
+    /// Creates a new `UniqueRc`.
     ///
     /// Weak references to this `UniqueRc` can be created with [`UniqueRc::downgrade`]. Upgrading
     /// these weak references will fail before the `UniqueRc` has been converted into an [`Rc`].
@@ -3569,54 +3580,78 @@ impl<T> UniqueRc<T> {
     #[cfg(not(no_global_oom_handling))]
     #[unstable(feature = "unique_rc_arc", issue = "112566")]
     pub fn new(value: T) -> Self {
-        Self {
-            ptr: Box::leak(Box::new(RcBox {
+        Self::new_in(value, Global)
+    }
+}
+
+impl<T, A: Allocator> UniqueRc<T, A> {
+    /// Creates a new `UniqueRc` in the provided allocator.
+    ///
+    /// Weak references to this `UniqueRc` can be created with [`UniqueRc::downgrade`]. Upgrading
+    /// these weak references will fail before the `UniqueRc` has been converted into an [`Rc`].
+    /// After converting the `UniqueRc` into an [`Rc`], any weak references created beforehand will
+    /// point to the new [`Rc`].
+    #[cfg(not(no_global_oom_handling))]
+    #[unstable(feature = "unique_rc_arc", issue = "112566")]
+    pub fn new_in(value: T, alloc: A) -> Self {
+        let (ptr, alloc) = Box::into_unique(Box::new_in(
+            RcBox {
                 strong: Cell::new(0),
                 // keep one weak reference so if all the weak pointers that are created are dropped
                 // the UniqueRc still stays valid.
                 weak: Cell::new(1),
                 value,
-            }))
-            .into(),
-            phantom: PhantomData,
-        }
-    }
-
-    /// Creates a new weak reference to the `UniqueRc`
-    ///
-    /// Attempting to upgrade this weak reference will fail before the `UniqueRc` has been converted
-    /// to a [`Rc`] using [`UniqueRc::into_rc`].
-    #[unstable(feature = "unique_rc_arc", issue = "112566")]
-    pub fn downgrade(this: &Self) -> Weak<T> {
-        // SAFETY: This pointer was allocated at creation time and we guarantee that we only have
-        // one strong reference before converting to a regular Rc.
-        unsafe {
-            this.ptr.as_ref().inc_weak();
-        }
-        Weak { ptr: this.ptr, alloc: Global }
+            },
+            alloc,
+        ));
+        Self { ptr: ptr.into(), phantom: PhantomData, alloc }
     }
+}
 
-    /// Converts the `UniqueRc` into a regular [`Rc`]
+impl<T: ?Sized, A: Allocator> UniqueRc<T, A> {
+    /// Converts the `UniqueRc` into a regular [`Rc`].
     ///
     /// This consumes the `UniqueRc` and returns a regular [`Rc`] that contains the `value` that
     /// is passed to `into_rc`.
     ///
     /// Any weak references created before this method is called can now be upgraded to strong
     /// references.
     #[unstable(feature = "unique_rc_arc", issue = "112566")]
-    pub fn into_rc(this: Self) -> Rc<T> {
+    pub fn into_rc(this: Self) -> Rc<T, A> {
         let mut this = ManuallyDrop::new(this);
+
+        // Move the allocator out.
+        // SAFETY: `this.alloc` will not be accessed again, nor dropped because it is in
+        // a `ManuallyDrop`.
+        let alloc: A = unsafe { ptr::read(&this.alloc) };
+
         // SAFETY: This pointer was allocated at creation time so we know it is valid.
         unsafe {
             // Convert our weak reference into a strong reference
             this.ptr.as_mut().strong.set(1);
-            Rc::from_inner(this.ptr)
+            Rc::from_inner_in(this.ptr, alloc)
+        }
+    }
+}
+
+impl<T: ?Sized, A: Allocator + Clone> UniqueRc<T, A> {
+    /// Creates a new weak reference to the `UniqueRc`.
+    ///
+    /// Attempting to upgrade this weak reference will fail before the `UniqueRc` has been converted
+    /// to a [`Rc`] using [`UniqueRc::into_rc`].
+    #[unstable(feature = "unique_rc_arc", issue = "112566")]
+    pub fn downgrade(this: &Self) -> Weak<T, A> {
+        // SAFETY: This pointer was allocated at creation time and we guarantee that we only have
+        // one strong reference before converting to a regular Rc.
+        unsafe {
+            this.ptr.as_ref().inc_weak();
         }
+        Weak { ptr: this.ptr, alloc: this.alloc.clone() }
     }
 }
 
 #[unstable(feature = "unique_rc_arc", issue = "112566")]
-impl<T> Deref for UniqueRc<T> {
+impl<T: ?Sized, A: Allocator> Deref for UniqueRc<T, A> {
     type Target = T;
 
     fn deref(&self) -> &T {
@@ -3626,7 +3661,7 @@ impl<T> Deref for UniqueRc<T> {
 }
 
 #[unstable(feature = "unique_rc_arc", issue = "112566")]
-impl<T> DerefMut for UniqueRc<T> {
+impl<T: ?Sized, A: Allocator> DerefMut for UniqueRc<T, A> {
     fn deref_mut(&mut self) -> &mut T {
         // SAFETY: This pointer was allocated at creation time so we know it is valid. We know we
         // have unique ownership and therefore it's safe to make a mutable reference because
@@ -3636,7 +3671,7 @@ impl<T> DerefMut for UniqueRc<T> {
 }
 
 #[unstable(feature = "unique_rc_arc", issue = "112566")]
-unsafe impl<#[may_dangle] T> Drop for UniqueRc<T> {
+unsafe impl<#[may_dangle] T: ?Sized, A: Allocator> Drop for UniqueRc<T, A> {
     fn drop(&mut self) {
         unsafe {
             // destroy the contained object
@@ -3646,7 +3681,7 @@ unsafe impl<#[may_dangle] T> Drop for UniqueRc<T> {
             self.ptr.as_ref().dec_weak();
 
             if self.ptr.as_ref().weak() == 0 {
-                Global.deallocate(self.ptr.cast(), Layout::for_value_raw(self.ptr.as_ptr()));
+                self.alloc.deallocate(self.ptr.cast(), Layout::for_value_raw(self.ptr.as_ptr()));
             }
         }
     }

diff --git a/library/alloc/src/rc/tests.rs b/library/alloc/src/rc/tests.rs
@@ -606,6 +606,23 @@ fn test_unique_rc_drops_contents() {
     assert!(dropped);
 }
 
+/// Exercise the non-default allocator usage.
+#[test]
+fn test_unique_rc_with_alloc_drops_contents() {
+    let mut dropped = false;
+    struct DropMe<'a>(&'a mut bool);
+    impl Drop for DropMe<'_> {
+        fn drop(&mut self) {
+            *self.0 = true;
+        }
+    }
+    {
+        let rc = UniqueRc::new_in(DropMe(&mut dropped), std::alloc::System);
+        drop(rc);
+    }
+    assert!(dropped);
+}
+
 #[test]
 fn test_unique_rc_weak_clone_holding_ref() {
     let mut v = UniqueRc::new(0u8);
@@ -614,3 +631,12 @@ fn test_unique_rc_weak_clone_holding_ref() {
     let _ = w.clone(); // touch weak count
     *r = 123;
 }
+
+#[test]
+fn test_unique_rc_unsizing_coercion() {
+    let mut rc: UniqueRc<[u8]> = UniqueRc::new([0u8; 3]);
+    assert_eq!(rc.len(), 3);
+    rc[0] = 123;
+    let rc: Rc<[u8]> = UniqueRc::into_rc(rc);
+    assert_eq!(*rc, [123, 0, 0]);
+}