From d19bed73f76af92f798f4705e4a9fabf95555c89 Mon Sep 17 00:00:00 2001 From: Ben Kimock Date: Sun, 16 Jun 2024 12:58:08 -0400 Subject: [PATCH 1/4] Add a precondition check for Layout::from_size_align_unchecked --- library/core/src/alloc/layout.rs | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/library/core/src/alloc/layout.rs b/library/core/src/alloc/layout.rs index 0b92767c93205..5f694616e59e1 100644 --- a/library/core/src/alloc/layout.rs +++ b/library/core/src/alloc/layout.rs @@ -4,6 +4,7 @@ // 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; @@ -118,6 +119,15 @@ impl Layout { #[inline] #[rustc_allow_const_fn_unstable(ptr_alignment_type)] pub const unsafe fn from_size_align_unchecked(size: usize, align: usize) -> Self { + assert_unsafe_precondition!( + check_library_ub, + "Layout::from_size_align_unchecked requires that align is a power of 2 \ + and the rounded-up allocation size does not exceed isize::MAX", + ( + size: usize = size, + align: usize = align, + ) => Layout::from_size_align(size, align).is_ok() + ); // SAFETY: the caller is required to uphold the preconditions. unsafe { Layout { size, align: Alignment::new_unchecked(align) } } } From b2ec2169992749a577cd3587934b008606dd7a17 Mon Sep 17 00:00:00 2001 From: Ben Kimock Date: Tue, 18 Jun 2024 14:30:29 -0400 Subject: [PATCH 2/4] Use Alignment::of more --- library/alloc/src/lib.rs | 1 + library/alloc/src/raw_vec.rs | 6 +++--- library/core/src/alloc/layout.rs | 4 +++- 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs index 895d1b8d59f2c..c122e14854f57 100644 --- a/library/alloc/src/lib.rs +++ b/library/alloc/src/lib.rs @@ -138,6 +138,7 @@ #![feature(maybe_uninit_uninit_array_transpose)] #![feature(panic_internals)] #![feature(pattern)] +#![feature(ptr_alignment_type)] #![feature(ptr_internals)] #![feature(ptr_metadata)] #![feature(ptr_sub_ptr)] diff --git a/library/alloc/src/raw_vec.rs b/library/alloc/src/raw_vec.rs index 1134c7f833e2b..569aac75d2aef 100644 --- a/library/alloc/src/raw_vec.rs +++ b/library/alloc/src/raw_vec.rs @@ -4,7 +4,7 @@ use core::alloc::LayoutError; use core::cmp; use core::hint; use core::mem::{self, ManuallyDrop, MaybeUninit, SizedTypeProperties}; -use core::ptr::{self, NonNull, Unique}; +use core::ptr::{self, Alignment, NonNull, Unique}; #[cfg(not(no_global_oom_handling))] use crate::alloc::handle_alloc_error; @@ -306,9 +306,9 @@ impl RawVec { // 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 align = Alignment::of::(); let size = mem::size_of::().unchecked_mul(self.cap.0); - let layout = Layout::from_size_align_unchecked(size, align); + let layout = Layout::from_size_alignment(size, align).unwrap_unchecked(); Some((self.ptr.cast().into(), layout)) } } diff --git a/library/core/src/alloc/layout.rs b/library/core/src/alloc/layout.rs index 5f694616e59e1..b577c531ff55c 100644 --- a/library/core/src/alloc/layout.rs +++ b/library/core/src/alloc/layout.rs @@ -97,8 +97,10 @@ impl Layout { } /// Internal helper constructor to skip revalidating alignment validity. + #[unstable(feature = "ptr_alignment_type", issue = "102070")] + #[rustc_const_unstable(feature = "ptr_alignment_type", issue = "102070")] #[inline] - const fn from_size_alignment(size: usize, align: Alignment) -> Result { + pub const fn from_size_alignment(size: usize, align: Alignment) -> Result { if size > Self::max_size_for_align(align) { return Err(LayoutError); } From 96bf871f0712bd189e74b1ddef614f7ae842472f Mon Sep 17 00:00:00 2001 From: Ben Kimock Date: Wed, 19 Jun 2024 16:31:30 -0400 Subject: [PATCH 3/4] Remove debug_assert from Result::unwrap_unchecked --- library/core/src/result.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/library/core/src/result.rs b/library/core/src/result.rs index f8cdcc000c50e..e5a235b2c9909 100644 --- a/library/core/src/result.rs +++ b/library/core/src/result.rs @@ -1481,7 +1481,6 @@ impl Result { #[track_caller] #[stable(feature = "option_result_unwrap_unchecked", since = "1.58.0")] pub unsafe fn unwrap_unchecked(self) -> T { - debug_assert!(self.is_ok()); match self { Ok(t) => t, // SAFETY: the safety contract must be upheld by the caller. @@ -1513,7 +1512,6 @@ impl Result { #[track_caller] #[stable(feature = "option_result_unwrap_unchecked", since = "1.58.0")] pub unsafe fn unwrap_err_unchecked(self) -> E { - debug_assert!(self.is_err()); match self { // SAFETY: the safety contract must be upheld by the caller. Ok(_) => unsafe { hint::unreachable_unchecked() }, From 44b1911dd064f1eb833547341d74ff5faca8dd60 Mon Sep 17 00:00:00 2001 From: Ben Kimock Date: Thu, 20 Jun 2024 15:23:27 -0400 Subject: [PATCH 4/4] Golf --- library/alloc/src/raw_vec.rs | 23 +++++++++++++++++------ library/core/src/alloc/layout.rs | 4 +--- library/core/src/result.rs | 2 ++ 3 files changed, 20 insertions(+), 9 deletions(-) diff --git a/library/alloc/src/raw_vec.rs b/library/alloc/src/raw_vec.rs index 569aac75d2aef..251c95827d6fd 100644 --- a/library/alloc/src/raw_vec.rs +++ b/library/alloc/src/raw_vec.rs @@ -4,7 +4,7 @@ use core::alloc::LayoutError; use core::cmp; use core::hint; use core::mem::{self, ManuallyDrop, MaybeUninit, SizedTypeProperties}; -use core::ptr::{self, Alignment, NonNull, Unique}; +use core::ptr::{self, NonNull, Unique}; #[cfg(not(no_global_oom_handling))] use crate::alloc::handle_alloc_error; @@ -297,7 +297,20 @@ impl RawVec { } fn current_memory(&self) -> Option<(NonNull, Layout)> { - if T::IS_ZST || self.cap.0 == 0 { + // Reduce the amount of code we need to monomorphize per `T`. + #[inline] + #[rustc_no_mir_inline] + unsafe fn inner(size: usize, align: usize, cap: usize) -> Layout { + // SAFETY: Precondition guaranteed by the caller + unsafe { + let size = size.unchecked_mul(cap); + Layout::from_size_align_unchecked(size, align) + } + } + + let cap = self.cap.0; + + if T::IS_ZST || cap == 0 { None } else { // We could use Layout::array here which ensures the absence of isize and usize overflows @@ -306,10 +319,8 @@ impl RawVec { // 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 = Alignment::of::(); - let size = mem::size_of::().unchecked_mul(self.cap.0); - let layout = Layout::from_size_alignment(size, align).unwrap_unchecked(); - Some((self.ptr.cast().into(), layout)) + let layout = inner(mem::size_of::(), mem::align_of::(), cap); + Some((self.non_null().cast(), layout)) } } } diff --git a/library/core/src/alloc/layout.rs b/library/core/src/alloc/layout.rs index b577c531ff55c..5f694616e59e1 100644 --- a/library/core/src/alloc/layout.rs +++ b/library/core/src/alloc/layout.rs @@ -97,10 +97,8 @@ impl Layout { } /// Internal helper constructor to skip revalidating alignment validity. - #[unstable(feature = "ptr_alignment_type", issue = "102070")] - #[rustc_const_unstable(feature = "ptr_alignment_type", issue = "102070")] #[inline] - pub const fn from_size_alignment(size: usize, align: Alignment) -> Result { + const fn from_size_alignment(size: usize, align: Alignment) -> Result { if size > Self::max_size_for_align(align) { return Err(LayoutError); } diff --git a/library/core/src/result.rs b/library/core/src/result.rs index e5a235b2c9909..f8cdcc000c50e 100644 --- a/library/core/src/result.rs +++ b/library/core/src/result.rs @@ -1481,6 +1481,7 @@ impl Result { #[track_caller] #[stable(feature = "option_result_unwrap_unchecked", since = "1.58.0")] pub unsafe fn unwrap_unchecked(self) -> T { + debug_assert!(self.is_ok()); match self { Ok(t) => t, // SAFETY: the safety contract must be upheld by the caller. @@ -1512,6 +1513,7 @@ impl Result { #[track_caller] #[stable(feature = "option_result_unwrap_unchecked", since = "1.58.0")] pub unsafe fn unwrap_err_unchecked(self) -> E { + debug_assert!(self.is_err()); match self { // SAFETY: the safety contract must be upheld by the caller. Ok(_) => unsafe { hint::unreachable_unchecked() },