diff --git a/library/alloc/src/rc.rs b/library/alloc/src/rc.rs index 006d813e5f9fa..f3cbfe27b3eed 100644 --- a/library/alloc/src/rc.rs +++ b/library/alloc/src/rc.rs @@ -293,6 +293,15 @@ struct RcBox { value: T, } +/// Calculate layout for `RcBox` 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)`, but this created a misaligned + // reference (see #54908). + Layout::new::>().extend(layout).unwrap().0.pad_to_align() +} + /// A single-threaded reference-counting pointer. 'Rc' stands for 'Reference /// Counted'. /// @@ -1334,11 +1343,7 @@ impl Rc { allocate: impl FnOnce(Layout) -> Result, AllocError>, mem_to_rcbox: impl FnOnce(*mut u8) -> *mut RcBox, ) -> *mut RcBox { - // Calculate layout using the given value layout. - // Previously, layout was calculated on the expression - // `&*(ptr as *const RcBox)`, but this created a misaligned - // reference (see #54908). - let layout = Layout::new::>().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 Rc { allocate: impl FnOnce(Layout) -> Result, AllocError>, mem_to_rcbox: impl FnOnce(*mut u8) -> *mut RcBox, ) -> Result<*mut RcBox, AllocError> { - // Calculate layout using the given value layout. - // Previously, layout was calculated on the expression - // `&*(ptr as *const RcBox)`, but this created a misaligned - // reference (see #54908). - let layout = Layout::new::>().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 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 Rc<[T]> { } } + /// Create an `Rc<[T]>` by reusing the underlying memory + /// of a `Vec`. 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) -> Result, Vec> { + let layout_elements = Layout::array::(v.len()).unwrap(); + let layout_allocation = Layout::array::(v.capacity()).unwrap(); + let layout_rcbox = rcbox_layout_for_value_layout(layout_elements); + let mut ptr = NonNull::new(v.as_mut_ptr()).expect("`Vec` stores `NonNull`"); + 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::(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::(), &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 From> for Rc<[T]> { /// assert_eq!(vec![1, 2, 3], *shared); /// ``` #[inline] - fn from(mut v: Vec) -> 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) -> 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 81cd770748854..37e07eb5998b3 100644 --- a/library/alloc/src/sync.rs +++ b/library/alloc/src/sync.rs @@ -333,6 +333,15 @@ struct ArcInner { data: T, } +/// Calculate layout for `ArcInner` 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)`, but this created a misaligned + // reference (see #54908). + Layout::new::>().extend(layout).unwrap().0.pad_to_align() +} + unsafe impl Send for ArcInner {} unsafe impl Sync for ArcInner {} @@ -1154,11 +1163,7 @@ impl Arc { allocate: impl FnOnce(Layout) -> Result, AllocError>, mem_to_arcinner: impl FnOnce(*mut u8) -> *mut ArcInner, ) -> *mut ArcInner { - // Calculate layout using the given value layout. - // Previously, layout was calculated on the expression - // `&*(ptr as *const ArcInner)`, but this created a misaligned - // reference (see #54908). - let layout = Layout::new::>().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 Arc { allocate: impl FnOnce(Layout) -> Result, AllocError>, mem_to_arcinner: impl FnOnce(*mut u8) -> *mut ArcInner, ) -> Result<*mut ArcInner, AllocError> { - // Calculate layout using the given value layout. - // Previously, layout was calculated on the expression - // `&*(ptr as *const ArcInner)`, but this created a misaligned - // reference (see #54908). - let layout = Layout::new::>().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 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 Arc<[T]> { } } + /// Create an `Arc<[T]>` by reusing the underlying memory + /// of a `Vec`. 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) -> Result, Vec> { + let layout_elements = Layout::array::(v.len()).unwrap(); + let layout_allocation = Layout::array::(v.capacity()).unwrap(); + let layout_arcinner = arcinner_layout_for_value_layout(layout_elements); + let mut ptr = NonNull::new(v.as_mut_ptr()).expect("`Vec` stores `NonNull`"); + 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::(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::(), &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 From> for Arc<[T]> { /// assert_eq!(&[1, 2, 3], &shared[..]); /// ``` #[inline] - fn from(mut v: Vec) -> 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) -> 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 ce40b5c9b0a0d..eb379e4d6a10f 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::(); + let arc: Arc<[_]> = v.into(); + unsafe { + assert_eq!( + arc.as_ptr().cast::().offset_from(addr), + (std::mem::size_of::() * 2) as isize, + "Vector allocation not reused" + ); + } +} diff --git a/library/alloc/tests/rc.rs b/library/alloc/tests/rc.rs index efb39a609665b..1d5f3c5200648 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::(); + let rc: Rc<[_]> = v.into(); + unsafe { + assert_eq!( + rc.as_ptr().cast::().offset_from(addr), + (std::mem::size_of::() * 2) as isize, + "Vector allocation not reused" + ); + } +}