Add {Rc, Arc}::make_mut_slice().

These functions behave identically to `make_mut()`, but operate on
`Arc<[T]>` instead of `Arc<T>`.

Author: kpreid

library/alloc/src/rc.rs

@@ -1773,6 +1773,91 @@ impl<T: Clone, A: Allocator + Clone> Rc<T, A> {
     }
 }
 
+impl<T: Clone, A: Allocator + Clone> Rc<[T], A> {
+    /// Makes a mutable reference into the given `Rc<[T]>`.
+    ///
+    /// This function has the same purpose and characteristics as [`make_mut`](Self::make_mut),
+    /// except that it applies to slices rather than [`Sized`] values.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(make_mut_slice)]
+    /// use std::rc::Rc;
+    ///
+    /// let mut data: Rc<[i32]> = Rc::new([10, 20, 30]);
+    ///
+    /// Rc::make_mut_slice(&mut data)[0] += 1;        // Won't clone anything
+    /// let mut other_data = Rc::clone(&data);        // Won't clone inner data
+    /// Rc::make_mut_slice(&mut data)[1] += 1;        // Clones inner data
+    /// Rc::make_mut_slice(&mut data)[2] += 1;        // Won't clone anything
+    /// Rc::make_mut_slice(&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]);
+    /// ```
+    ///
+    /// [`Weak`] pointers will be disassociated:
+    ///
+    /// ```
+    /// #![feature(make_mut_slice)]
+    /// use std::rc::Rc;
+    ///
+    /// let mut data: Rc<[i32]> = Rc::new([75]);
+    /// let weak = Rc::downgrade(&data);
+    ///
+    /// assert!([75] == *data);
+    /// assert!([75] == *weak.upgrade().unwrap());
+    ///
+    /// Rc::make_mut_slice(&mut data)[0] += 1;
+    ///
+    /// assert!([76] == *data);
+    /// assert!(weak.upgrade().is_none());
+    /// ```
+    #[cfg(not(no_global_oom_handling))]
+    #[inline]
+    #[unstable(feature = "make_mut_slice", issue = "none")]
+    pub fn make_mut_slice(this: &mut Self) -> &mut [T] {
+        // This code is identical to `make_mut()` except that it clones len() items instead of 1.
+
+        let len = this.len();
+
+        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_slice_in(len, this.alloc.clone());
+            unsafe {
+                let data = Rc::get_mut_unchecked(&mut rc);
+                for i in 0..len {
+                    this[i].write_clone_into_raw(data[i].as_mut_ptr());
+                }
+                *this = rc.assume_init();
+            }
+        } else if Rc::weak_count(this) != 0 {
+            // Can just steal the data, all that's left is Weaks
+            let mut rc = Self::new_uninit_slice_in(len, this.alloc.clone());
+            unsafe {
+                let data = Rc::get_mut_unchecked(&mut rc);
+                mem::MaybeUninit::slice_as_mut_ptr(data)
+                    .copy_from_nonoverlapping(this.as_ptr(), len);
+
+                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());
+            }
+        }
+        // This unsafety is ok because we're guaranteed that the pointer
+        // returned is the *only* pointer that will ever be returned to T. Our
+        // reference count is guaranteed to be 1 at this point, and we required
+        // the `Rc<T>` itself to be `mut`, so we're returning the only possible
+        // reference to the allocation.
+        unsafe { &mut this.ptr.as_mut().value }
+    }
+}
+
 impl<A: Allocator + Clone> Rc<dyn Any, A> {
     /// Attempt to downcast the `Rc<dyn Any>` to a concrete type.
     ///

library/alloc/src/sync.rs

@@ -2198,6 +2198,112 @@ impl<T: Clone, A: Allocator + Clone> Arc<T, A> {
     }
 }
 
+impl<T: Clone, A: Allocator + Clone> Arc<[T], A> {
+    /// Makes a mutable reference into the given `Arc<[T]>`.
+    ///
+    /// This function has the same purpose and characteristics as [`make_mut`](Self::make_mut),
+    /// except that it applies to slices rather than [`Sized`] values.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(make_mut_slice)]
+    /// use std::sync::Arc;
+    ///
+    /// let mut data: Arc<[i32]> = Rc::new([10, 20, 30]);
+    ///
+    /// Arc::make_mut_slice(&mut data)[0] += 1;        // Won't clone anything
+    /// let mut other_data = Arc::clone(&data);        // Won't clone inner data
+    /// Arc::make_mut_slice(&mut data)[1] += 1;        // Clones inner data
+    /// Arc::make_mut_slice(&mut data)[2] += 1;        // Won't clone anything
+    /// Arc::make_mut_slice(&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]);
+    /// ```
+    ///
+    /// [`Weak`] pointers will be disassociated:
+    ///
+    /// ```
+    /// #![feature(make_mut_slice)]
+    /// use std::sync::Arc;
+    ///
+    /// let mut data: Arc<[i32]> = Arc::new([75]);
+    /// let weak = Arc::downgrade(&data);
+    ///
+    /// assert!([75] == *data);
+    /// assert!([75] == *weak.upgrade().unwrap());
+    ///
+    /// Arc::make_mut_slice(&mut data)[0] += 1;
+    ///
+    /// assert!([76] == *data);
+    /// assert!(weak.upgrade().is_none());
+    /// ```
+    #[cfg(not(no_global_oom_handling))]
+    #[inline]
+    #[stable(feature = "arc_unique", since = "1.4.0")]
+    pub fn make_mut_slice(this: &mut Self) -> &mut [T] {
+        // This code is identical to `make_mut()` except that it clones len() items instead of 1.
+
+        let len = this.len();
+
+        // 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.
+        //
+        // Use Acquire to ensure that we see any writes to `weak` that happen
+        // before release writes (i.e., decrements) to `strong`. Since we hold a
+        // weak count, there's no chance the ArcInner itself could be
+        // 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_slice_in(len, this.alloc.clone());
+            unsafe {
+                let data = Arc::get_mut_unchecked(&mut arc);
+                for i in 0..len {
+                    this[i].write_clone_into_raw(data[i].as_mut_ptr());
+                }
+                *this = arc.assume_init();
+            }
+        } 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
+            // dropped. Worst case, we end up allocated a new Arc unnecessarily.
+
+            // We removed the last strong ref, but there are additional weak
+            // refs remaining. We'll move the contents to a new Arc, and
+            // invalidate the other weak refs.
+
+            // Note that it is not possible for the read of `weak` to yield
+            // usize::MAX (i.e., locked), since the weak count can only be
+            // locked by a thread with a strong reference.
+
+            // Materialize our own implicit weak pointer, so that it can clean
+            // up the ArcInner as needed.
+            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_slice_in(len, this.alloc.clone());
+            unsafe {
+                let data = Arc::get_mut_unchecked(&mut arc);
+                mem::MaybeUninit::slice_as_mut_ptr(data)
+                    .copy_from_nonoverlapping(this.as_ptr(), len);
+                ptr::write(this, arc.assume_init());
+            }
+        } else {
+            // We were the sole reference of either kind; bump back up the
+            // strong ref count.
+            this.inner().strong.store(1, Release);
+        }
+
+        // As with `get_mut()`, the unsafety is ok because our reference was
+        // either unique to begin with, or became one upon cloning the contents.
+        unsafe { Self::get_mut_unchecked(this) }
+    }
+}
+
 impl<T: ?Sized, A: Allocator> Arc<T, A> {
     /// Returns a mutable reference into the given `Arc`, if there are
     /// no other `Arc` or [`Weak`] pointers to the same allocation.