Skip to content

Commit

Permalink
Do not allocate for ZST ThinBox attempt 2 (using const_allocate)
Browse files Browse the repository at this point in the history
There's PR #123184
which avoids allocation for ZST ThinBox.

That PR has an issue with unsoundness with misuse of `MaybeUninit`
(see comments in that PR).

This PR is much simpler implementation which does not have this
problem, but it uses `const_allocate` feature.
  • Loading branch information
stepancheg committed Mar 31, 2024
1 parent 204805a commit 9546c3f
Show file tree
Hide file tree
Showing 2 changed files with 74 additions and 13 deletions.
85 changes: 72 additions & 13 deletions library/alloc/src/boxed/thin.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,14 @@
use crate::alloc::{self, Layout, LayoutError};
use core::error::Error;
use core::fmt::{self, Debug, Display, Formatter};
#[cfg(not(no_global_oom_handling))]
use core::intrinsics::const_allocate;
use core::marker::PhantomData;
#[cfg(not(no_global_oom_handling))]
use core::marker::Unsize;
use core::mem::{self, SizedTypeProperties};
use core::mem;
#[cfg(not(no_global_oom_handling))]
use core::mem::SizedTypeProperties;
use core::ops::{Deref, DerefMut};
use core::ptr::Pointee;
use core::ptr::{self, NonNull};
Expand Down Expand Up @@ -91,6 +95,58 @@ impl<T> ThinBox<T> {

#[unstable(feature = "thin_box", issue = "92791")]
impl<Dyn: ?Sized> ThinBox<Dyn> {
#[cfg(not(no_global_oom_handling))]
fn new_unsize_zst<T>(value: T) -> Self
where
T: Unsize<Dyn>,
{
assert!(mem::size_of::<T>() == 0);

const fn max(a: usize, b: usize) -> usize {
if a > b { a } else { b }
}

let alloc: &<Dyn as Pointee>::Metadata = const {
// ZST T pointer must be aligned to both the value and the metadata.
// Metadata must be at the address `&T - size_of::<Metadata>`.
// This is how we need to allocate the memory:
// [ padding | metadata | ZST T value ]

// FIXME: just call `WithHeader::alloc_layout` with size reset to 0.
// Currently that's blocked on `Layout::extend` not being `const fn`.

let alloc_align =
max(mem::align_of::<T>(), mem::align_of::<<Dyn as Pointee>::Metadata>());

let alloc_size =
max(mem::align_of::<T>(), mem::size_of::<<Dyn as Pointee>::Metadata>());

unsafe {
// SAFETY: align is power of two because it is the maximum of two alignments.
let alloc: *mut u8 = const_allocate(alloc_size, alloc_align);

let metadata_offset =
alloc_size.checked_sub(mem::size_of::<<Dyn as Pointee>::Metadata>()).unwrap();
// SAFETY: adding offset within the allocation.
let metadata_ptr: *mut <Dyn as Pointee>::Metadata =
alloc.add(metadata_offset).cast();
// SAFETY: `*metadata_ptr` is within the allocation.
metadata_ptr.write(ptr::metadata::<Dyn>(ptr::dangling::<T>() as *const Dyn));

// SAFETY: we have just written the metadata.
&*(metadata_ptr)
}
};

// SAFETY: `alloc` points to `<Dyn as Pointee>::Metadata`, so addition does not overflow.
let value_ptr =
unsafe { (alloc as *const <Dyn as Pointee>::Metadata).add(1) }.cast::<u8>().cast_mut();
let ptr = WithOpaqueHeader(NonNull::new(value_ptr).unwrap());
// Forget the value to avoid double drop.
mem::forget(value);
ThinBox::<Dyn> { ptr, _marker: PhantomData }
}

/// Moves a type to the heap with its [`Metadata`] stored in the heap allocation instead of on
/// the stack.
///
Expand All @@ -109,9 +165,13 @@ impl<Dyn: ?Sized> ThinBox<Dyn> {
where
T: Unsize<Dyn>,
{
let meta = ptr::metadata(&value as &Dyn);
let ptr = WithOpaqueHeader::new(meta, value);
ThinBox { ptr, _marker: PhantomData }
if mem::size_of::<T>() == 0 {
Self::new_unsize_zst(value)
} else {
let meta = ptr::metadata(&value as &Dyn);
let ptr = WithOpaqueHeader::new(meta, value);
ThinBox { ptr, _marker: PhantomData }
}
}
}

Expand Down Expand Up @@ -300,20 +360,19 @@ impl<H> WithHeader<H> {

impl<H> Drop for DropGuard<H> {
fn drop(&mut self) {
// All ZST are allocated statically.
if self.value_layout.size() == 0 {
return;
}

unsafe {
// SAFETY: Layout must have been computable if we're in drop
let (layout, value_offset) =
WithHeader::<H>::alloc_layout(self.value_layout).unwrap_unchecked();

// Note: Don't deallocate if the layout size is zero, because the pointer
// didn't come from the allocator.
if layout.size() != 0 {
alloc::dealloc(self.ptr.as_ptr().sub(value_offset), layout);
} else {
debug_assert!(
value_offset == 0 && H::IS_ZST && self.value_layout.size() == 0
);
}
// Since we only allocate for non-ZSTs, the layout size cannot be zero.
debug_assert!(layout.size() != 0);
alloc::dealloc(self.ptr.as_ptr().sub(value_offset), layout);
}
}
}
Expand Down
2 changes: 2 additions & 0 deletions library/alloc/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -114,8 +114,10 @@
#![feature(const_box)]
#![feature(const_cow_is_borrowed)]
#![feature(const_eval_select)]
#![feature(const_heap)]
#![feature(const_maybe_uninit_as_mut_ptr)]
#![feature(const_maybe_uninit_write)]
#![feature(const_option)]
#![feature(const_pin)]
#![feature(const_refs_to_cell)]
#![feature(const_size_of_val)]
Expand Down

0 comments on commit 9546c3f

Please sign in to comment.