Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 2 commits into from
Nov 17, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
84 changes: 65 additions & 19 deletions library/alloc/src/rc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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'.
///
Expand Down Expand Up @@ -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))
Expand All @@ -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)?;
Expand Down Expand Up @@ -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))]
Expand All @@ -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.
Expand Down Expand Up @@ -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
}
}
}
}
}
Expand Down
85 changes: 66 additions & 19 deletions library/alloc/src/sync.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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> {}

Expand Down Expand Up @@ -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))
Expand All @@ -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)?;

Expand Down Expand Up @@ -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))]
Expand All @@ -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.
Expand Down Expand Up @@ -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
}
}
}
}
}
Expand Down
15 changes: 15 additions & 0 deletions library/alloc/tests/arc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think there is UB in this test: #104565.

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"
);
}
}
15 changes: 15 additions & 0 deletions library/alloc/tests/rc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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"
);
}
}