Rollup merge of #90066 - yaahc:thinbox, r=joshtriplett

Add new ThinBox type for 1 stack pointer wide heap allocated trait objects

**Zulip Thread**: https://rust-lang.zulipchat.com/#narrow/stream/219381-t-libs/topic/ThinBox

Based on https://github.com/matthieu-m/rfc2580/blob/b58d1d3cba0d4b5e859d3617ea2d0943aaa31329/examples/thin.rs

Tracking Issue: #92791

Usage Trial: https://github.com/yaahc/pgx/pull/1/files

## TODO

- [x] make sure to test with #[repr(align(1024))] structs etc

library/alloc/src/boxed.rs

@@ -163,6 +163,11 @@ use crate::str::from_boxed_utf8_unchecked;
 #[cfg(not(no_global_oom_handling))]
 use crate::vec::Vec;
 
+#[unstable(feature = "thin_box", issue = "92791")]
+pub use thin::ThinBox;
+
+mod thin;
+
 /// A pointer type for heap allocation.
 ///
 /// See the [module-level documentation](../../std/boxed/index.html) for more.

library/alloc/src/boxed/thin.rs

@@ -0,0 +1,215 @@
+// Based on
+// https://github.com/matthieu-m/rfc2580/blob/b58d1d3cba0d4b5e859d3617ea2d0943aaa31329/examples/thin.rs
+// by matthieu-m
+use crate::alloc::{self, Layout, LayoutError};
+use core::fmt::{self, Debug, Display, Formatter};
+use core::marker::{PhantomData, Unsize};
+use core::mem;
+use core::ops::{Deref, DerefMut};
+use core::ptr::Pointee;
+use core::ptr::{self, NonNull};
+
+/// ThinBox.
+///
+/// A thin pointer for heap allocation, regardless of T.
+///
+/// # Examples
+///
+/// ```
+/// #![feature(thin_box)]
+/// use std::boxed::ThinBox;
+///
+/// let five = ThinBox::new(5);
+/// let thin_slice = ThinBox::<[i32]>::new_unsize([1, 2, 3, 4]);
+///
+/// use std::mem::{size_of, size_of_val};
+/// let size_of_ptr = size_of::<*const ()>();
+/// assert_eq!(size_of_ptr, size_of_val(&five));
+/// assert_eq!(size_of_ptr, size_of_val(&thin_slice));
+/// ```
+#[unstable(feature = "thin_box", issue = "92791")]
+pub struct ThinBox<T: ?Sized> {
+    ptr: WithHeader<<T as Pointee>::Metadata>,
+    _marker: PhantomData<T>,
+}
+
+#[unstable(feature = "thin_box", issue = "92791")]
+impl<T> ThinBox<T> {
+    /// Moves a type to the heap with its `Metadata` stored in the heap allocation instead of on
+    /// the stack.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(thin_box)]
+    /// use std::boxed::ThinBox;
+    ///
+    /// let five = ThinBox::new(5);
+    /// ```
+    #[cfg(not(no_global_oom_handling))]
+    pub fn new(value: T) -> Self {
+        let meta = ptr::metadata(&value);
+        let ptr = WithHeader::new(meta, value);
+        ThinBox { ptr, _marker: PhantomData }
+    }
+}
+
+#[unstable(feature = "thin_box", issue = "92791")]
+impl<Dyn: ?Sized> ThinBox<Dyn> {
+    /// Moves a type to the heap with its `Metadata` stored in the heap allocation instead of on
+    /// the stack.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(thin_box)]
+    /// use std::boxed::ThinBox;
+    ///
+    /// let thin_slice = ThinBox::<[i32]>::new_unsize([1, 2, 3, 4]);
+    /// ```
+    #[cfg(not(no_global_oom_handling))]
+    pub fn new_unsize<T>(value: T) -> Self
+    where
+        T: Unsize<Dyn>,
+    {
+        let meta = ptr::metadata(&value as &Dyn);
+        let ptr = WithHeader::new(meta, value);
+        ThinBox { ptr, _marker: PhantomData }
+    }
+}
+
+#[unstable(feature = "thin_box", issue = "92791")]
+impl<T: ?Sized + Debug> Debug for ThinBox<T> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        Debug::fmt(self.deref(), f)
+    }
+}
+
+#[unstable(feature = "thin_box", issue = "92791")]
+impl<T: ?Sized + Display> Display for ThinBox<T> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        Display::fmt(self.deref(), f)
+    }
+}
+
+#[unstable(feature = "thin_box", issue = "92791")]
+impl<T: ?Sized> Deref for ThinBox<T> {
+    type Target = T;
+
+    fn deref(&self) -> &T {
+        let value = self.data();
+        let metadata = self.meta();
+        let pointer = ptr::from_raw_parts(value as *const (), metadata);
+        unsafe { &*pointer }
+    }
+}
+
+#[unstable(feature = "thin_box", issue = "92791")]
+impl<T: ?Sized> DerefMut for ThinBox<T> {
+    fn deref_mut(&mut self) -> &mut T {
+        let value = self.data();
+        let metadata = self.meta();
+        let pointer = ptr::from_raw_parts_mut::<T>(value as *mut (), metadata);
+        unsafe { &mut *pointer }
+    }
+}
+
+#[unstable(feature = "thin_box", issue = "92791")]
+impl<T: ?Sized> Drop for ThinBox<T> {
+    fn drop(&mut self) {
+        unsafe {
+            let value = self.deref_mut();
+            let value = value as *mut T;
+            self.ptr.drop::<T>(value);
+        }
+    }
+}
+
+#[unstable(feature = "thin_box", issue = "92791")]
+impl<T: ?Sized> ThinBox<T> {
+    fn meta(&self) -> <T as Pointee>::Metadata {
+        //  Safety:
+        //  -   NonNull and valid.
+        unsafe { *self.ptr.header() }
+    }
+
+    fn data(&self) -> *mut u8 {
+        self.ptr.value()
+    }
+}
+
+/// A pointer to type-erased data, guaranteed to have a header `H` before the pointed-to location.
+struct WithHeader<H>(NonNull<u8>, PhantomData<H>);
+
+impl<H> WithHeader<H> {
+    #[cfg(not(no_global_oom_handling))]
+    fn new<T>(header: H, value: T) -> WithHeader<H> {
+        let value_layout = Layout::new::<T>();
+        let Ok((layout, value_offset)) = Self::alloc_layout(value_layout) else {
+            // We pass an empty layout here because we do not know which layout caused the
+            // arithmetic overflow in `Layout::extend` and `handle_alloc_error` takes `Layout` as
+            // its argument rather than `Result<Layout, LayoutError>`, also this function has been
+            // stable since 1.28 ._.
+            //
+            // On the other hand, look at this gorgeous turbofish!
+            alloc::handle_alloc_error(Layout::new::<()>());
+        };
+
+        unsafe {
+            let ptr = alloc::alloc(layout);
+
+            if ptr.is_null() {
+                alloc::handle_alloc_error(layout);
+            }
+            //  Safety:
+            //  -   The size is at least `aligned_header_size`.
+            let ptr = ptr.add(value_offset) as *mut _;
+
+            let ptr = NonNull::new_unchecked(ptr);
+
+            let result = WithHeader(ptr, PhantomData);
+            ptr::write(result.header(), header);
+            ptr::write(result.value().cast(), value);
+
+            result
+        }
+    }
+
+    //  Safety:
+    //  -   Assumes that `value` can be dereferenced.
+    unsafe fn drop<T: ?Sized>(&self, value: *mut T) {
+        unsafe {
+            // SAFETY: Layout must have been computable if we're in drop
+            let (layout, value_offset) =
+                Self::alloc_layout(Layout::for_value_raw(value)).unwrap_unchecked();
+
+            ptr::drop_in_place::<T>(value);
+            // We only drop the value because the Pointee trait requires that the metadata is copy
+            // aka trivially droppable
+            alloc::dealloc(self.0.as_ptr().sub(value_offset), layout);
+        }
+    }
+
+    fn header(&self) -> *mut H {
+        //  Safety:
+        //  - At least `size_of::<H>()` bytes are allocated ahead of the pointer.
+        //  - We know that H will be aligned because the middle pointer is aligned to the greater
+        //    of the alignment of the header and the data and the header size includes the padding
+        //    needed to align the header. Subtracting the header size from the aligned data pointer
+        //    will always result in an aligned header pointer, it just may not point to the
+        //    beginning of the allocation.
+        unsafe { self.0.as_ptr().sub(Self::header_size()) as *mut H }
+    }
+
+    fn value(&self) -> *mut u8 {
+        self.0.as_ptr()
+    }
+
+    const fn header_size() -> usize {
+        mem::size_of::<H>()
+    }
+
+    fn alloc_layout(value_layout: Layout) -> Result<(Layout, usize), LayoutError> {
+        Layout::new::<H>().extend(value_layout)
+    }
+}

library/alloc/src/lib.rs

@@ -120,6 +120,7 @@
 #![feature(nonnull_slice_from_raw_parts)]
 #![feature(pattern)]
 #![feature(ptr_internals)]
+#![feature(ptr_metadata)]
 #![feature(receiver_trait)]
 #![feature(set_ptr_value)]
 #![feature(slice_group_by)]
@@ -152,6 +153,7 @@
 #![feature(fundamental)]
 #![cfg_attr(not(test), feature(generator_trait))]
 #![feature(lang_items)]
+#![feature(let_else)]
 #![feature(min_specialization)]
 #![feature(negative_impls)]
 #![feature(never_type)]

library/alloc/tests/lib.rs

@@ -39,6 +39,7 @@
 #![feature(nonnull_slice_from_raw_parts)]
 #![feature(panic_update_hook)]
 #![feature(slice_flatten)]
+#![feature(thin_box)]
 
 use std::collections::hash_map::DefaultHasher;
 use std::hash::{Hash, Hasher};
@@ -57,6 +58,7 @@ mod rc;
 mod slice;
 mod str;
 mod string;
+mod thin_box;
 mod vec;
 mod vec_deque;
 

library/alloc/tests/thin_box.rs

@@ -0,0 +1,26 @@
+use alloc::boxed::ThinBox;
+use core::mem::size_of;
+
+#[test]
+fn want_niche_optimization() {
+    fn uses_niche<T: ?Sized>() -> bool {
+        size_of::<*const ()>() == size_of::<Option<ThinBox<T>>>()
+    }
+
+    trait Tr {}
+    assert!(uses_niche::<dyn Tr>());
+    assert!(uses_niche::<[i32]>());
+    assert!(uses_niche::<i32>());
+}
+
+#[test]
+fn want_thin() {
+    fn is_thin<T: ?Sized>() -> bool {
+        size_of::<*const ()>() == size_of::<ThinBox<T>>()
+    }
+
+    trait Tr {}
+    assert!(is_thin::<dyn Tr>());
+    assert!(is_thin::<[i32]>());
+    assert!(is_thin::<i32>());
+}

library/std/src/error.rs

@@ -516,6 +516,14 @@ impl<T: Error> Error for Box<T> {
     }
 }
 
+#[unstable(feature = "thin_box", issue = "92791")]
+impl<T: ?Sized + crate::error::Error> crate::error::Error for crate::boxed::ThinBox<T> {
+    fn source(&self) -> Option<&(dyn crate::error::Error + 'static)> {
+        use core::ops::Deref;
+        self.deref().source()
+    }
+}
+
 #[stable(feature = "error_by_ref", since = "1.51.0")]
 impl<'a, T: Error + ?Sized> Error for &'a T {
     #[allow(deprecated, deprecated_in_future)]

library/std/src/lib.rs

@@ -290,6 +290,7 @@
 #![feature(get_mut_unchecked)]
 #![feature(map_try_insert)]
 #![feature(new_uninit)]
+#![feature(thin_box)]
 #![feature(toowned_clone_into)]
 #![feature(try_reserve_kind)]
 #![feature(vec_into_raw_parts)]

src/test/ui/box/thin_align.rs

@@ -0,0 +1,26 @@
+#![feature(thin_box)]
+// run-pass
+use std::boxed::ThinBox;
+use std::error::Error;
+use std::ops::Deref;
+use std::fmt;
+
+fn main() {
+    let expected = "Foo error!";
+    let a: ThinBox<dyn Error> = ThinBox::new_unsize(Foo(expected));
+    let a = a.deref();
+    let msg = a.to_string();
+    assert_eq!(expected, msg);
+}
+
+#[derive(Debug)]
+#[repr(align(1024))]
+struct Foo(&'static str);
+
+impl fmt::Display for Foo {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
+impl Error for Foo {}

src/test/ui/box/thin_drop.rs

@@ -0,0 +1,37 @@
+#![feature(thin_box)]
+// run-pass
+use std::boxed::ThinBox;
+use std::error::Error;
+use std::ops::Deref;
+use std::fmt;
+
+fn main() {
+    let expected = "Foo error!";
+    let mut dropped = false;
+    {
+        let foo = Foo(expected, &mut dropped);
+        let a: ThinBox<dyn Error> = ThinBox::new_unsize(foo);
+        let a = a.deref();
+        let msg = a.to_string();
+        assert_eq!(expected, msg);
+    }
+    assert!(dropped);
+}
+
+#[derive(Debug)]
+#[repr(align(1024))]
+struct Foo<'a>(&'static str, &'a mut bool);
+
+impl Drop for Foo<'_> {
+    fn drop(&mut self) {
+        *self.1 = true;
+    }
+}
+
+impl fmt::Display for Foo<'_> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
+impl Error for Foo<'_> {}