Skip to content

Commit

Permalink
Do not allocate for ZST ThinBox
Browse files Browse the repository at this point in the history
  • Loading branch information
stepancheg committed Mar 29, 2024
1 parent db2f975 commit 2048773
Show file tree
Hide file tree
Showing 2 changed files with 96 additions and 13 deletions.
107 changes: 94 additions & 13 deletions library/alloc/src/boxed/thin.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,10 @@
use crate::alloc::{self, Layout, LayoutError};
use core::error::Error;
use core::fmt::{self, Debug, Display, Formatter};
use core::marker::PhantomData;
#[cfg(not(no_global_oom_handling))]
use core::marker::Unsize;
use core::mem::{self, SizedTypeProperties};
use core::marker::{Freeze, PhantomData};
use core::mem::{self, MaybeUninit, SizedTypeProperties};
use core::ops::{Deref, DerefMut};
use core::ptr::Pointee;
use core::ptr::{self, NonNull};
Expand Down Expand Up @@ -91,6 +91,83 @@ 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>,
{
#[repr(C)]
struct ReprC<A, B> {
a: A,
b: B,
}

// Allocate header like this:
// ```
// [ ... | header ]
// ```
// where the struct is aligned to both header and value.
#[repr(C)]
struct AlignedHeader<H: Copy, T> {
header_data: MaybeUninit<ReprC<H, [T; 0]>>,
}

impl<H: Copy, T> AlignedHeader<H, T> {
const fn make(header: H) -> Self {
let mut data = MaybeUninit::<ReprC<H, [T; 0]>>::zeroed();
unsafe {
data.as_mut_ptr().add(1).cast::<H>().sub(1).write(header);
}
AlignedHeader { header_data: data }
}
}

#[repr(C)]
struct DynZstAlloc<T, Dyn: ?Sized> {
header: AlignedHeader<<Dyn as Pointee>::Metadata, T>,
value: MaybeUninit<T>,
}

// We need `Freeze` so we could call `&DynZstAlloc::LAYOUT`.
// SAFETY: data is immutable.
unsafe impl<T, Dyn: ?Sized> Freeze for DynZstAlloc<T, Dyn> {}

impl<T, Dyn: ?Sized> DynZstAlloc<T, Dyn>
where
T: Unsize<Dyn>,
{
const LAYOUT: DynZstAlloc<T, Dyn> = DynZstAlloc {
header: AlignedHeader::make(ptr::metadata::<Dyn>(
ptr::dangling::<T>() as *const Dyn
)),
value: MaybeUninit::uninit(),
};

fn static_alloc<'a>() -> &'a DynZstAlloc<T, Dyn> {
&Self::LAYOUT
}
}

let alloc: &DynZstAlloc<T, Dyn> = DynZstAlloc::<T, Dyn>::static_alloc();

let value_offset = mem::offset_of!(DynZstAlloc<T, Dyn>, value);
assert_eq!(value_offset, mem::size_of::<AlignedHeader<<Dyn as Pointee>::Metadata, T>>());

let ptr = WithOpaqueHeader(
NonNull::new(
// SAFETY: there's no overflow here because we add field offset.
unsafe {
(alloc as *const DynZstAlloc<T, Dyn> as *mut u8).add(value_offset) as *mut _
},
)
.unwrap(),
);
let thin_box = ThinBox::<Dyn> { ptr, _marker: PhantomData };
// Forget the value to avoid double drop.
mem::forget(value);
thin_box
}

/// Moves a type to the heap with its [`Metadata`] stored in the heap allocation instead of on
/// the stack.
///
Expand All @@ -109,9 +186,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 && mem::size_of::<<Dyn as Pointee>::Metadata>() != 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 @@ -155,7 +236,7 @@ impl<T: ?Sized> DerefMut for ThinBox<T> {
impl<T: ?Sized> Drop for ThinBox<T> {
fn drop(&mut self) {
unsafe {
let value = self.deref_mut();
let value: &mut T = self.deref_mut();
let value = value as *mut T;
self.with_header().drop::<T>(value);
}
Expand Down Expand Up @@ -300,20 +381,20 @@ 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
);
}
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 @@ -130,6 +130,8 @@
#![feature(extend_one)]
#![feature(fmt_internals)]
#![feature(fn_traits)]
#![feature(freeze)]
#![feature(freeze_impls)]
#![feature(generic_nonzero)]
#![feature(hasher_prefixfree_extras)]
#![feature(hint_assert_unchecked)]
Expand Down

0 comments on commit 2048773

Please sign in to comment.