Generalize {Rc,Arc}::make_mut() to unsized types.

kpreid · kpreid · commit 8ebd67508195 · 2023-09-24T17:53:49.000-07:00
This requires introducing a new internal type `RcUninit`, which can own
an `RcBox&lt;T&gt;` without requiring it to be initialized, sized, or a slice.
diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs
@@ -141,6 +141,7 @@
 #![feature(maybe_uninit_uninit_array_transpose)]
 #![feature(pattern)]
 #![feature(pointer_byte_offsets)]
+#![feature(ptr_from_ref)]
 #![feature(ptr_internals)]
 #![feature(ptr_metadata)]
 #![feature(ptr_sub_ptr)]
diff --git a/library/alloc/src/rc.rs b/library/alloc/src/rc.rs
@@ -258,8 +258,7 @@ use core::intrinsics::abort;
 use core::iter;
 use core::marker::{PhantomData, Unsize};
 #[cfg(not(no_global_oom_handling))]
-use core::mem::size_of_val;
-use core::mem::{self, align_of_val_raw, forget, ManuallyDrop};
+use core::mem::{self, align_of_val_raw, forget, size_of_val, ManuallyDrop};
 use core::ops::{CoerceUnsized, Deref, DerefMut, DispatchFromDyn, Receiver};
 use core::panic::{RefUnwindSafe, UnwindSafe};
 #[cfg(not(no_global_oom_handling))]
@@ -1651,7 +1650,7 @@ impl<T: ?Sized, A: Allocator> Rc<T, A> {
     }
 }
 
-impl<T: Clone, A: Allocator + Clone> Rc<T, A> {
+impl<T: ?Sized + CloneToUninit, A: Allocator + Clone> Rc<T, A> {
     /// Makes a mutable reference into the given `Rc`.
     ///
     /// If there are other `Rc` pointers to the same allocation, then `make_mut` will
@@ -1706,27 +1705,47 @@ impl<T: Clone, A: Allocator + Clone> Rc<T, A> {
     #[inline]
     #[stable(feature = "rc_unique", since = "1.4.0")]
     pub fn make_mut(this: &mut Self) -> &mut T {
+        let size_of_val = mem::size_of_val::<T>(&**this);
+
         if Rc::strong_count(this) != 1 {
             // Gotta clone the data, there are other Rcs.
-            // Pre-allocate memory to allow writing the cloned value directly.
-            let mut rc = Self::new_uninit_in(this.alloc.clone());
-            unsafe {
-                let data = Rc::get_mut_unchecked(&mut rc);
-                (**this).clone_to_uninit(data.as_mut_ptr());
-                *this = rc.assume_init();
-            }
+
+            let this_data_ref: &T = &**this;
+            // `in_progress` drops the allocation if we panic before finishing initializing it.
+            let mut in_progress: RcUninit<T, A> = RcUninit::new(this_data_ref, this.alloc.clone());
+
+            // Initialize with clone of this.
+            let initialized_clone = unsafe {
+                // Clone. If the clone panics, `in_progress` will be dropped and clean up.
+                this_data_ref.clone_to_uninit(in_progress.data_ptr());
+                // Cast type of pointer, now that it is initialized.
+                in_progress.into_rc()
+            };
+
+            // Replace `this` with newly constructed Rc.
+            *this = initialized_clone;
         } else if Rc::weak_count(this) != 0 {
             // Can just steal the data, all that's left is Weaks
-            let mut rc = Self::new_uninit_in(this.alloc.clone());
+
+            // We don't need panic-protection like the above branch does, but we might as well
+            // use the same mechanism.
+            let mut in_progress: RcUninit<T, A> = RcUninit::new(&**this, this.alloc.clone());
             unsafe {
-                let data = Rc::get_mut_unchecked(&mut rc);
-                data.as_mut_ptr().copy_from_nonoverlapping(&**this, 1);
+                // Initialize `in_progress` with move of **this.
+                // We have to express this in terms of bytes because `T: ?Sized`; there is no
+                // operation that just copies a value based on its `size_of_val()`.
+                ptr::copy_nonoverlapping(
+                    ptr::from_ref(&**this).cast::<u8>(),
+                    in_progress.data_ptr().cast::<u8>(),
+                    size_of_val,
+                );
 
                 this.inner().dec_strong();
                 // Remove implicit strong-weak ref (no need to craft a fake
                 // Weak here -- we know other Weaks can clean up for us)
                 this.inner().dec_weak();
-                ptr::write(this, rc.assume_init());
+                // Replace `this` with newly constructed Rc that has the moved data.
+                ptr::write(this, in_progress.into_rc());
             }
         }
         // This unsafety is ok because we're guaranteed that the pointer
@@ -1736,7 +1755,9 @@ impl<T: Clone, A: Allocator + Clone> Rc<T, A> {
         // reference to the allocation.
         unsafe { &mut this.ptr.as_mut().value }
     }
+}
 
+impl<T: Clone, A: Allocator + Clone> Rc<T, A> {
     /// If we have the only reference to `T` then unwrap it. Otherwise, clone `T` and return the
     /// clone.
     ///
@@ -3360,6 +3381,66 @@ fn data_offset_align(align: usize) -> usize {
     layout.size() + layout.padding_needed_for(align)
 }
 
+/// A unique owning pointer to a [`RcBox`] **that does not imply the contents are initialized,**
+/// but will deallocate it (without dropping the value) when dropped.
+///
+/// This is a helper for [`Rc::make_mut()`] to ensure correct cleanup on panic.
+struct RcUninit<T: ?Sized, A: Allocator> {
+    ptr: NonNull<RcBox<T>>,
+    layout_for_value: Layout,
+    alloc: Option<A>,
+}
+
+impl<T: ?Sized, A: Allocator> RcUninit<T, A> {
+    /// Allocate a RcBox with layout suitable to contain `for_value` or a clone of it.
+    #[cfg(not(no_global_oom_handling))]
+    fn new(for_value: &T, alloc: A) -> RcUninit<T, A> {
+        let layout = Layout::for_value(for_value);
+        let ptr = unsafe {
+            Rc::allocate_for_layout(
+                layout,
+                |layout_for_rcbox| alloc.allocate(layout_for_rcbox),
+                |mem| mem.with_metadata_of(ptr::from_ref(for_value) as *const RcBox<T>),
+            )
+        };
+        Self { ptr: NonNull::new(ptr).unwrap(), layout_for_value: layout, alloc: Some(alloc) }
+    }
+
+    /// Returns the pointer to be written into to initialize the [`Rc`].
+    fn data_ptr(&mut self) -> *mut T {
+        let offset = data_offset_align(self.layout_for_value.align());
+        unsafe { self.ptr.as_ptr().byte_add(offset) as *mut T }
+    }
+
+    /// Upgrade this into a normal [`Rc`].
+    ///
+    /// # Safety
+    ///
+    /// The data must have been initialized (by writing to [`Self::data_ptr()`]).
+    unsafe fn into_rc(mut self) -> Rc<T, A> {
+        let ptr = self.ptr;
+        let alloc = self.alloc.take().unwrap();
+        mem::forget(self);
+        // SAFETY: The pointer is valid as per `RcUninit::new`, and the caller is responsible
+        // for having initialized the data.
+        unsafe { Rc::from_ptr_in(ptr.as_ptr(), alloc) }
+    }
+}
+
+impl<T: ?Sized, A: Allocator> Drop for RcUninit<T, A> {
+    fn drop(&mut self) {
+        // SAFETY:
+        // * new() produced a pointer safe to deallocate.
+        // * We own the pointer unless into_rc() was called, which forgets us.
+        unsafe {
+            self.alloc
+                .take()
+                .unwrap()
+                .deallocate(self.ptr.cast(), rcbox_layout_for_value_layout(self.layout_for_value));
+        }
+    }
+}
+
 /// A uniquely owned `Rc`
 ///
 /// This represents an `Rc` that is known to be uniquely owned -- that is, have exactly one strong
diff --git a/library/alloc/src/rc/tests.rs b/library/alloc/src/rc/tests.rs
@@ -321,6 +321,24 @@ fn test_cowrc_clone_weak() {
     assert!(cow1_weak.upgrade().is_none());
 }
 
+/// This is similar to the doc-test for `Rc::make_mut()`, but on an unsized type (slice).
+#[test]
+fn test_cowrc_unsized() {
+    use std::rc::Rc;
+
+    let mut data: Rc<[i32]> = Rc::new([10, 20, 30]);
+
+    Rc::make_mut(&mut data)[0] += 1; // Won't clone anything
+    let mut other_data = Rc::clone(&data); // Won't clone inner data
+    Rc::make_mut(&mut data)[1] += 1; // Clones inner data
+    Rc::make_mut(&mut data)[2] += 1; // Won't clone anything
+    Rc::make_mut(&mut other_data)[0] *= 10; // Won't clone anything
+
+    // Now `data` and `other_data` point to different allocations.
+    assert_eq!(*data, [11, 21, 31]);
+    assert_eq!(*other_data, [110, 20, 30]);
+}
+
 #[test]
 fn test_show() {
     let foo = Rc::new(75);
diff --git a/library/alloc/src/sync.rs b/library/alloc/src/sync.rs
@@ -2055,7 +2055,7 @@ impl<T: ?Sized, A: Allocator> Deref for Arc<T, A> {
 #[unstable(feature = "receiver_trait", issue = "none")]
 impl<T: ?Sized> Receiver for Arc<T> {}
 
-impl<T: Clone, A: Allocator + Clone> Arc<T, A> {
+impl<T: ?Sized + CloneToUninit, A: Allocator + Clone> Arc<T, A> {
     /// Makes a mutable reference into the given `Arc`.
     ///
     /// If there are other `Arc` pointers to the same allocation, then `make_mut` will
@@ -2110,6 +2110,8 @@ impl<T: Clone, A: Allocator + Clone> Arc<T, A> {
     #[inline]
     #[stable(feature = "arc_unique", since = "1.4.0")]
     pub fn make_mut(this: &mut Self) -> &mut T {
+        let size_of_val = mem::size_of_val::<T>(&**this);
+
         // Note that we hold both a strong reference and a weak reference.
         // Thus, releasing our strong reference only will not, by itself, cause
         // the memory to be deallocated.
@@ -2120,13 +2122,19 @@ impl<T: Clone, A: Allocator + Clone> Arc<T, A> {
         // deallocated.
         if this.inner().strong.compare_exchange(1, 0, Acquire, Relaxed).is_err() {
             // Another strong pointer exists, so we must clone.
-            // Pre-allocate memory to allow writing the cloned value directly.
-            let mut arc = Self::new_uninit_in(this.alloc.clone());
-            unsafe {
-                let data = Arc::get_mut_unchecked(&mut arc);
-                (**this).clone_to_uninit(data.as_mut_ptr());
-                *this = arc.assume_init();
-            }
+
+            let this_data_ref: &T = &**this;
+            // `in_progress` drops the allocation if we panic before finishing initializing it.
+            let mut in_progress: ArcUninit<T, A> =
+                ArcUninit::new(this_data_ref, this.alloc.clone());
+
+            let initialized_clone = unsafe {
+                // Clone. If the clone panics, `in_progress` will be dropped and clean up.
+                this_data_ref.clone_to_uninit(in_progress.data_ptr());
+                // Cast type of pointer, now that it is initialized.
+                in_progress.into_arc()
+            };
+            *this = initialized_clone;
         } else if this.inner().weak.load(Relaxed) != 1 {
             // Relaxed suffices in the above because this is fundamentally an
             // optimization: we are always racing with weak pointers being
@@ -2145,11 +2153,21 @@ impl<T: Clone, A: Allocator + Clone> Arc<T, A> {
             let _weak = Weak { ptr: this.ptr, alloc: this.alloc.clone() };
 
             // Can just steal the data, all that's left is Weaks
-            let mut arc = Self::new_uninit_in(this.alloc.clone());
+            //
+            // We don't need panic-protection like the above branch does, but we might as well
+            // use the same mechanism.
+            let mut in_progress: ArcUninit<T, A> = ArcUninit::new(&**this, this.alloc.clone());
             unsafe {
-                let data = Arc::get_mut_unchecked(&mut arc);
-                data.as_mut_ptr().copy_from_nonoverlapping(&**this, 1);
-                ptr::write(this, arc.assume_init());
+                // Initialize `in_progress` with move of **this.
+                // We have to express this in terms of bytes because `T: ?Sized`; there is no
+                // operation that just copies a value based on its `size_of_val()`.
+                ptr::copy_nonoverlapping(
+                    ptr::from_ref(&**this).cast::<u8>(),
+                    in_progress.data_ptr().cast::<u8>(),
+                    size_of_val,
+                );
+
+                ptr::write(this, in_progress.into_arc());
             }
         } else {
             // We were the sole reference of either kind; bump back up the
@@ -2161,7 +2179,9 @@ impl<T: Clone, A: Allocator + Clone> Arc<T, A> {
         // either unique to begin with, or became one upon cloning the contents.
         unsafe { Self::get_mut_unchecked(this) }
     }
+}
 
+impl<T: Clone, A: Allocator + Clone> Arc<T, A> {
     /// If we have the only reference to `T` then unwrap it. Otherwise, clone `T` and return the
     /// clone.
     ///
@@ -3557,6 +3577,66 @@ fn data_offset_align(align: usize) -> usize {
     layout.size() + layout.padding_needed_for(align)
 }
 
+/// A unique owning pointer to a [`ArcInner`] **that does not imply the contents are initialized,**
+/// but will deallocate it (without dropping the value) when dropped.
+///
+/// This is a helper for [`Arc::make_mut()`] to ensure correct cleanup on panic.
+struct ArcUninit<T: ?Sized, A: Allocator> {
+    ptr: NonNull<ArcInner<T>>,
+    layout_for_value: Layout,
+    alloc: Option<A>,
+}
+
+impl<T: ?Sized, A: Allocator> ArcUninit<T, A> {
+    /// Allocate a ArcInner with layout suitable to contain `for_value` or a clone of it.
+    #[cfg(not(no_global_oom_handling))]
+    fn new(for_value: &T, alloc: A) -> ArcUninit<T, A> {
+        let layout = Layout::for_value(for_value);
+        let ptr = unsafe {
+            Arc::allocate_for_layout(
+                layout,
+                |layout_for_arcinner| alloc.allocate(layout_for_arcinner),
+                |mem| mem.with_metadata_of(ptr::from_ref(for_value) as *const ArcInner<T>),
+            )
+        };
+        Self { ptr: NonNull::new(ptr).unwrap(), layout_for_value: layout, alloc: Some(alloc) }
+    }
+
+    /// Returns the pointer to be written into to initialize the [`Arc`].
+    fn data_ptr(&mut self) -> *mut T {
+        let offset = data_offset_align(self.layout_for_value.align());
+        unsafe { self.ptr.as_ptr().byte_add(offset) as *mut T }
+    }
+
+    /// Upgrade this into a normal [`Arc`].
+    ///
+    /// # Safety
+    ///
+    /// The data must have been initialized (by writing to [`Self::data_ptr()`]).
+    unsafe fn into_arc(mut self) -> Arc<T, A> {
+        let ptr = self.ptr;
+        let alloc = self.alloc.take().unwrap();
+        mem::forget(self);
+        // SAFETY: The pointer is valid as per `ArcUninit::new`, and the caller is responsible
+        // for having initialized the data.
+        unsafe { Arc::from_ptr_in(ptr.as_ptr(), alloc) }
+    }
+}
+
+impl<T: ?Sized, A: Allocator> Drop for ArcUninit<T, A> {
+    fn drop(&mut self) {
+        // SAFETY:
+        // * new() produced a pointer safe to deallocate.
+        // * We own the pointer unless into_arc() was called, which forgets us.
+        unsafe {
+            self.alloc.take().unwrap().deallocate(
+                self.ptr.cast(),
+                arcinner_layout_for_value_layout(self.layout_for_value),
+            );
+        }
+    }
+}
+
 #[stable(feature = "arc_error", since = "1.52.0")]
 impl<T: core::error::Error + ?Sized> core::error::Error for Arc<T> {
     #[allow(deprecated, deprecated_in_future)]
diff --git a/library/alloc/tests/arc.rs b/library/alloc/tests/arc.rs
@@ -210,3 +210,21 @@ fn weak_may_dangle() {
     // `val` dropped here while still borrowed
     // borrow might be used here, when `val` is dropped and runs the `Drop` code for type `std::sync::Weak`
 }
+
+/// This is similar to the doc-test for `Arc::make_mut()`, but on an unsized type (slice).
+#[test]
+fn make_mut_unsized() {
+    use alloc::sync::Arc;
+
+    let mut data: Arc<[i32]> = Arc::new([10, 20, 30]);
+
+    Arc::make_mut(&mut data)[0] += 1; // Won't clone anything
+    let mut other_data = Arc::clone(&data); // Won't clone inner data
+    Arc::make_mut(&mut data)[1] += 1; // Clones inner data
+    Arc::make_mut(&mut data)[2] += 1; // Won't clone anything
+    Arc::make_mut(&mut other_data)[0] *= 10; // Won't clone anything
+
+    // Now `data` and `other_data` point to different allocations.
+    assert_eq!(*data, [11, 21, 31]);
+    assert_eq!(*other_data, [110, 20, 30]);
+}
