Auto merge of #104205 - clubby789:grow-rc, r=thomcc

Attempt to reuse `Vec<T>` backing storage for `Rc/Arc<[T]>`

If a `Vec<T>` has sufficient capacity to store the inner `RcBox<[T]>`, we can just reuse the existing allocation and shift the elements up, instead of making a new allocation.

4 files changed, 161 insertions, 38 deletions

diff --git a/library/alloc/src/rc.rs b/library/alloc/src/rc.rs
index 006d813e5f9..f3cbfe27b3e 100644
--- a/library/alloc/src/rc.rs
+++ b/library/alloc/src/rc.rs
@@ -293,6 +293,15 @@ struct RcBox<T: ?Sized> {
     value: T,
 }
 
+/// Calculate layout for `RcBox<T>` using the inner value's layout
+fn rcbox_layout_for_value_layout(layout: Layout) -> Layout {
+    // Calculate layout using the given value layout.
+    // Previously, layout was calculated on the expression
+    // `&*(ptr as *const RcBox<T>)`, but this created a misaligned
+    // reference (see #54908).
+    Layout::new::<RcBox<()>>().extend(layout).unwrap().0.pad_to_align()
+}
+
 /// A single-threaded reference-counting pointer. 'Rc' stands for 'Reference
 /// Counted'.
 ///
@@ -1334,11 +1343,7 @@ impl<T: ?Sized> Rc<T> {
         allocate: impl FnOnce(Layout) -> Result<NonNull<[u8]>, AllocError>,
         mem_to_rcbox: impl FnOnce(*mut u8) -> *mut RcBox<T>,
     ) -> *mut RcBox<T> {
-        // Calculate layout using the given value layout.
-        // Previously, layout was calculated on the expression
-        // `&*(ptr as *const RcBox<T>)`, but this created a misaligned
-        // reference (see #54908).
-        let layout = Layout::new::<RcBox<()>>().extend(value_layout).unwrap().0.pad_to_align();
+        let layout = rcbox_layout_for_value_layout(value_layout);
         unsafe {
             Rc::try_allocate_for_layout(value_layout, allocate, mem_to_rcbox)
                 .unwrap_or_else(|_| handle_alloc_error(layout))
@@ -1357,11 +1362,7 @@ impl<T: ?Sized> Rc<T> {
         allocate: impl FnOnce(Layout) -> Result<NonNull<[u8]>, AllocError>,
         mem_to_rcbox: impl FnOnce(*mut u8) -> *mut RcBox<T>,
     ) -> Result<*mut RcBox<T>, AllocError> {
-        // Calculate layout using the given value layout.
-        // Previously, layout was calculated on the expression
-        // `&*(ptr as *const RcBox<T>)`, but this created a misaligned
-        // reference (see #54908).
-        let layout = Layout::new::<RcBox<()>>().extend(value_layout).unwrap().0.pad_to_align();
+        let layout = rcbox_layout_for_value_layout(value_layout);
 
         // Allocate for the layout.
         let ptr = allocate(layout)?;
@@ -1428,7 +1429,7 @@ impl<T> Rc<[T]> {
         }
     }
 
-    /// Copy elements from slice into newly allocated Rc<\[T\]>
+    /// Copy elements from slice into newly allocated `Rc<[T]>`
     ///
     /// Unsafe because the caller must either take ownership or bind `T: Copy`
     #[cfg(not(no_global_oom_handling))]
@@ -1440,6 +1441,48 @@ impl<T> Rc<[T]> {
         }
     }
 
+    /// Create an `Rc<[T]>` by reusing the underlying memory
+    /// of a `Vec<T>`. This will return the vector if the existing allocation
+    /// is not large enough.
+    #[cfg(not(no_global_oom_handling))]
+    fn try_from_vec_in_place(mut v: Vec<T>) -> Result<Rc<[T]>, Vec<T>> {
+        let layout_elements = Layout::array::<T>(v.len()).unwrap();
+        let layout_allocation = Layout::array::<T>(v.capacity()).unwrap();
+        let layout_rcbox = rcbox_layout_for_value_layout(layout_elements);
+        let mut ptr = NonNull::new(v.as_mut_ptr()).expect("`Vec<T>` stores `NonNull<T>`");
+        if layout_rcbox.size() > layout_allocation.size()
+            || layout_rcbox.align() > layout_allocation.align()
+        {
+            // Can't fit - calling `grow` would involve `realloc`
+            // (which copies the elements), followed by copying again.
+            return Err(v);
+        }
+        if layout_rcbox.size() < layout_allocation.size()
+            || layout_rcbox.align() < layout_allocation.align()
+        {
+            // We need to shrink the allocation so that it fits
+            // https://doc.rust-lang.org/nightly/std/alloc/trait.Allocator.html#memory-fitting
+            // SAFETY:
+            // - Vec allocates by requesting `Layout::array::<T>(capacity)`, so this capacity matches
+            // - `layout_rcbox` is smaller
+            // If this fails, the ownership has not been transferred
+            if let Ok(p) = unsafe { Global.shrink(ptr.cast(), layout_allocation, layout_rcbox) } {
+                ptr = p.cast();
+            } else {
+                return Err(v);
+            }
+        }
+        // Make sure the vec's memory isn't deallocated now
+        let v = mem::ManuallyDrop::new(v);
+        let ptr: *mut RcBox<[T]> = ptr::slice_from_raw_parts_mut(ptr.as_ptr(), v.len()) as _;
+        unsafe {
+            ptr::copy(ptr.cast::<T>(), &mut (*ptr).value as *mut [T] as *mut T, v.len());
+            ptr::write(&mut (*ptr).strong, Cell::new(1));
+            ptr::write(&mut (*ptr).weak, Cell::new(1));
+            Ok(Self::from_ptr(ptr))
+        }
+    }
+
     /// Constructs an `Rc<[T]>` from an iterator known to be of a certain size.
     ///
     /// Behavior is undefined should the size be wrong.
@@ -1965,14 +2008,17 @@ impl<T> From<Vec<T>> for Rc<[T]> {
     /// assert_eq!(vec![1, 2, 3], *shared);
     /// ```
     #[inline]
-    fn from(mut v: Vec<T>) -> Rc<[T]> {
-        unsafe {
-            let rc = Rc::copy_from_slice(&v);
-
-            // Allow the Vec to free its memory, but not destroy its contents
-            v.set_len(0);
-
-            rc
+    fn from(v: Vec<T>) -> Rc<[T]> {
+        match Rc::try_from_vec_in_place(v) {
+            Ok(rc) => rc,
+            Err(mut v) => {
+                unsafe {
+                    let rc = Rc::copy_from_slice(&v);
+                    // Allow the Vec to free its memory, but not destroy its contents
+                    v.set_len(0);
+                    rc
+                }
+            }
         }
     }
 }
diff --git a/library/alloc/src/sync.rs b/library/alloc/src/sync.rs
index 81cd7707488..37e07eb5998 100644
--- a/library/alloc/src/sync.rs
+++ b/library/alloc/src/sync.rs
@@ -333,6 +333,15 @@ struct ArcInner<T: ?Sized> {
     data: T,
 }
 
+/// Calculate layout for `ArcInner<T>` using the inner value's layout
+fn arcinner_layout_for_value_layout(layout: Layout) -> Layout {
+    // Calculate layout using the given value layout.
+    // Previously, layout was calculated on the expression
+    // `&*(ptr as *const ArcInner<T>)`, but this created a misaligned
+    // reference (see #54908).
+    Layout::new::<ArcInner<()>>().extend(layout).unwrap().0.pad_to_align()
+}
+
 unsafe impl<T: ?Sized + Sync + Send> Send for ArcInner<T> {}
 unsafe impl<T: ?Sized + Sync + Send> Sync for ArcInner<T> {}
 
@@ -1154,11 +1163,7 @@ impl<T: ?Sized> Arc<T> {
         allocate: impl FnOnce(Layout) -> Result<NonNull<[u8]>, AllocError>,
         mem_to_arcinner: impl FnOnce(*mut u8) -> *mut ArcInner<T>,
     ) -> *mut ArcInner<T> {
-        // Calculate layout using the given value layout.
-        // Previously, layout was calculated on the expression
-        // `&*(ptr as *const ArcInner<T>)`, but this created a misaligned
-        // reference (see #54908).
-        let layout = Layout::new::<ArcInner<()>>().extend(value_layout).unwrap().0.pad_to_align();
+        let layout = arcinner_layout_for_value_layout(value_layout);
         unsafe {
             Arc::try_allocate_for_layout(value_layout, allocate, mem_to_arcinner)
                 .unwrap_or_else(|_| handle_alloc_error(layout))
@@ -1176,11 +1181,7 @@ impl<T: ?Sized> Arc<T> {
         allocate: impl FnOnce(Layout) -> Result<NonNull<[u8]>, AllocError>,
         mem_to_arcinner: impl FnOnce(*mut u8) -> *mut ArcInner<T>,
     ) -> Result<*mut ArcInner<T>, AllocError> {
-        // Calculate layout using the given value layout.
-        // Previously, layout was calculated on the expression
-        // `&*(ptr as *const ArcInner<T>)`, but this created a misaligned
-        // reference (see #54908).
-        let layout = Layout::new::<ArcInner<()>>().extend(value_layout).unwrap().0.pad_to_align();
+        let layout = arcinner_layout_for_value_layout(value_layout);
 
         let ptr = allocate(layout)?;
 
@@ -1246,7 +1247,7 @@ impl<T> Arc<[T]> {
         }
     }
 
-    /// Copy elements from slice into newly allocated Arc<\[T\]>
+    /// Copy elements from slice into newly allocated `Arc<[T]>`
     ///
     /// Unsafe because the caller must either take ownership or bind `T: Copy`.
     #[cfg(not(no_global_oom_handling))]
@@ -1260,6 +1261,49 @@ impl<T> Arc<[T]> {
         }
     }
 
+    /// Create an `Arc<[T]>` by reusing the underlying memory
+    /// of a `Vec<T>`. This will return the vector if the existing allocation
+    /// is not large enough.
+    #[cfg(not(no_global_oom_handling))]
+    fn try_from_vec_in_place(mut v: Vec<T>) -> Result<Arc<[T]>, Vec<T>> {
+        let layout_elements = Layout::array::<T>(v.len()).unwrap();
+        let layout_allocation = Layout::array::<T>(v.capacity()).unwrap();
+        let layout_arcinner = arcinner_layout_for_value_layout(layout_elements);
+        let mut ptr = NonNull::new(v.as_mut_ptr()).expect("`Vec<T>` stores `NonNull<T>`");
+        if layout_arcinner.size() > layout_allocation.size()
+            || layout_arcinner.align() > layout_allocation.align()
+        {
+            // Can't fit - calling `grow` would involve `realloc`
+            // (which copies the elements), followed by copying again.
+            return Err(v);
+        }
+        if layout_arcinner.size() < layout_allocation.size()
+            || layout_arcinner.align() < layout_allocation.align()
+        {
+            // We need to shrink the allocation so that it fits
+            // https://doc.rust-lang.org/nightly/std/alloc/trait.Allocator.html#memory-fitting
+            // SAFETY:
+            // - Vec allocates by requesting `Layout::array::<T>(capacity)`, so this capacity matches
+            // - `layout_arcinner` is smaller
+            // If this fails, the ownership has not been transferred
+            if let Ok(p) = unsafe { Global.shrink(ptr.cast(), layout_allocation, layout_arcinner) }
+            {
+                ptr = p.cast();
+            } else {
+                return Err(v);
+            }
+        }
+        // Make sure the vec's memory isn't deallocated now
+        let v = mem::ManuallyDrop::new(v);
+        let ptr: *mut ArcInner<[T]> = ptr::slice_from_raw_parts_mut(ptr.as_ptr(), v.len()) as _;
+        unsafe {
+            ptr::copy(ptr.cast::<T>(), &mut (*ptr).data as *mut [T] as *mut T, v.len());
+            ptr::write(&mut (*ptr).strong, atomic::AtomicUsize::new(1));
+            ptr::write(&mut (*ptr).weak, atomic::AtomicUsize::new(1));
+            Ok(Self::from_ptr(ptr))
+        }
+    }
+
     /// Constructs an `Arc<[T]>` from an iterator known to be of a certain size.
     ///
     /// Behavior is undefined should the size be wrong.
@@ -2571,14 +2615,17 @@ impl<T> From<Vec<T>> for Arc<[T]> {
     /// assert_eq!(&[1, 2, 3], &shared[..]);
     /// ```
     #[inline]
-    fn from(mut v: Vec<T>) -> Arc<[T]> {
-        unsafe {
-            let arc = Arc::copy_from_slice(&v);
-
-            // Allow the Vec to free its memory, but not destroy its contents
-            v.set_len(0);
-
-            arc
+    fn from(v: Vec<T>) -> Arc<[T]> {
+        match Arc::try_from_vec_in_place(v) {
+            Ok(rc) => rc,
+            Err(mut v) => {
+                unsafe {
+                    let rc = Arc::copy_from_slice(&v);
+                    // Allow the Vec to free its memory, but not destroy its contents
+                    v.set_len(0);
+                    rc
+                }
+            }
         }
     }
 }
diff --git a/library/alloc/tests/arc.rs b/library/alloc/tests/arc.rs
index ce40b5c9b0a..eb379e4d6a1 100644
--- a/library/alloc/tests/arc.rs
+++ b/library/alloc/tests/arc.rs
@@ -210,3 +210,18 @@ 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`
 }
+
+#[test]
+fn arc_from_vec_opt() {
+    let mut v = Vec::with_capacity(64);
+    v.push(0usize);
+    let addr = v.as_ptr().cast::<u8>();
+    let arc: Arc<[_]> = v.into();
+    unsafe {
+        assert_eq!(
+            arc.as_ptr().cast::<u8>().offset_from(addr),
+            (std::mem::size_of::<usize>() * 2) as isize,
+            "Vector allocation not reused"
+        );
+    }
+}
diff --git a/library/alloc/tests/rc.rs b/library/alloc/tests/rc.rs
index efb39a60966..1d5f3c52006 100644
--- a/library/alloc/tests/rc.rs
+++ b/library/alloc/tests/rc.rs
@@ -206,3 +206,18 @@ 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::rc::Weak`
 }
+
+#[test]
+fn rc_from_vec_opt() {
+    let mut v = Vec::with_capacity(64);
+    v.push(0usize);
+    let addr = v.as_ptr().cast::<u8>();
+    let rc: Rc<[_]> = v.into();
+    unsafe {
+        assert_eq!(
+            rc.as_ptr().cast::<u8>().offset_from(addr),
+            (std::mem::size_of::<usize>() * 2) as isize,
+            "Vector allocation not reused"
+        );
+    }
+}