From 1b3be3ece2ed3a722b5cad2898a17a8f71b997dc Mon Sep 17 00:00:00 2001 From: Scott McMurray Date: Sun, 23 Jun 2024 11:34:17 -0700 Subject: [PATCH] Add a `Layout::array_unchecked` for RawVec to use Instead of it writing the code for it itself. --- library/alloc/src/raw_vec.rs | 14 ++----- library/core/src/alloc/layout.rs | 67 ++++++++++++++++++++++++-------- 2 files changed, 54 insertions(+), 27 deletions(-) diff --git a/library/alloc/src/raw_vec.rs b/library/alloc/src/raw_vec.rs index 7b7dae5a057f0..788f1ff11d674 100644 --- a/library/alloc/src/raw_vec.rs +++ b/library/alloc/src/raw_vec.rs @@ -300,17 +300,9 @@ impl RawVec { if T::IS_ZST || self.cap.0 == 0 { None } else { - // We could use Layout::array here which ensures the absence of isize and usize overflows - // and could hypothetically handle differences between stride and size, but this memory - // has already been allocated so we know it can't overflow and currently Rust does not - // support such types. So we can do better by skipping some checks and avoid an unwrap. - const { assert!(mem::size_of::() % mem::align_of::() == 0) }; - unsafe { - let align = mem::align_of::(); - let size = mem::size_of::().unchecked_mul(self.cap.0); - let layout = Layout::from_size_align_unchecked(size, align); - Some((self.ptr.cast().into(), layout)) - } + // SAFETY: This memory has already been allocated so we know it can't overflow + let layout = unsafe { Layout::array_unchecked::(self.cap.0) }; + Some((self.ptr.cast().into(), layout)) } } diff --git a/library/core/src/alloc/layout.rs b/library/core/src/alloc/layout.rs index 0b92767c93205..c0cc97d923254 100644 --- a/library/core/src/alloc/layout.rs +++ b/library/core/src/alloc/layout.rs @@ -4,9 +4,11 @@ // collections, resulting in having to optimize down excess IR multiple times. // Your performance intuition is useless. Run perf. +use crate::assert_unsafe_precondition; use crate::cmp; use crate::error::Error; use crate::fmt; +use crate::intrinsics; use crate::mem; use crate::ptr::{Alignment, NonNull}; @@ -425,43 +427,76 @@ impl Layout { Layout::from_size_alignment(new_size, self.align) } - /// Creates a layout describing the record for a `[T; n]`. + /// Creates a layout describing the record for a `[T; len]`. /// /// On arithmetic overflow or when the total size would exceed /// `isize::MAX`, returns `LayoutError`. #[stable(feature = "alloc_layout_manipulation", since = "1.44.0")] #[rustc_const_unstable(feature = "const_alloc_layout", issue = "67521")] #[inline] - pub const fn array(n: usize) -> Result { + pub const fn array(len: usize) -> Result { // Reduce the amount of code we need to monomorphize per `T`. - return inner(mem::size_of::(), Alignment::of::(), n); + return inner(Layout::new::(), len); #[inline] - const fn inner( - element_size: usize, - align: Alignment, - n: usize, - ) -> Result { + const fn inner(element: Layout, len: usize) -> Result { // We need to check two things about the size: // - That the total size won't overflow a `usize`, and // - That the total size still fits in an `isize`. // By using division we can check them both with a single threshold. // That'd usually be a bad idea, but thankfully here the element size // and alignment are constants, so the compiler will fold all of it. - if element_size != 0 && n > Layout::max_size_for_align(align) / element_size { + if element.size != 0 && len > Layout::max_size_for_align(element.align) / element.size { return Err(LayoutError); } - // SAFETY: We just checked that we won't overflow `usize` when we multiply. + // SAFETY: We just checked above that the `array_size` will not + // exceed `isize::MAX` even when rounded up to the alignment. + unsafe { Ok(Layout::array_of_unchecked(element, len)) } + } + } + + /// Creates a layout describing the record for a `[T; n]`, without checking + /// whether the `len` is an allowed length. + /// + /// # Safety + /// + /// `len` must be small enough that the total size *in bytes* of the + /// resulting layout is `isize::MAX` or less. + /// + /// Equivalently, either `T` is a ZST or `len <= isize::MAX / size_of::()`. + #[unstable(feature = "alloc_layout_extra", issue = "55724")] + #[rustc_const_unstable(feature = "const_alloc_layout", issue = "67521")] + #[inline] + pub const unsafe fn array_unchecked(len: usize) -> Self { + // SAFETY: `T` is a Rust type, so its size is a multiple of align, + // and our safety precondition ensures that len is small enough. + unsafe { Self::array_of_unchecked(Layout::new::(), len) } + } + + /// # Safety + /// - The element layout must be that of a rust type, with size a multiple of align + /// - The total size of the resulting array must be at most `isize::MAX` + #[inline] + const unsafe fn array_of_unchecked(element: Layout, len: usize) -> Layout { + assert_unsafe_precondition!( + check_language_ub, + "array_of_unchecked element must have size an integer multiple of \ + alignment and the total length must be be isize::MAX or less", + ( + size: usize = element.size, + align: usize = element.align.as_usize(), + len: usize = len, + ) => size % align == 0 && (size as u128) * (len as u128) <= (isize::MAX as u128), + ); + + // SAFETY: Out preconditions are exactly what's needed here + unsafe { // This is a useless hint inside this function, but after inlining this helps // deduplicate checks for whether the overall capacity is zero (e.g., in RawVec's // allocation path) before/after this multiplication. - let array_size = unsafe { element_size.unchecked_mul(n) }; - - // SAFETY: We just checked above that the `array_size` will not - // exceed `isize::MAX` even when rounded up to the alignment. - // And `Alignment` guarantees it's a power of two. - unsafe { Ok(Layout::from_size_align_unchecked(array_size, align.as_usize())) } + let bytes = intrinsics::unchecked_mul(element.size, len); + Self { size: bytes, align: element.align } } } }