From 7f5d282185674cac26f9f916f0f1ae8ec04e3482 Mon Sep 17 00:00:00 2001
From: Ben Kimock <kimockb@gmail.com>
Date: Sun, 16 Jun 2024 12:58:08 -0400
Subject: [PATCH 1/2] Add a precondition check for
 Layout::from_size_align_unchecked

---
 library/core/src/alloc/layout.rs | 32 +++++++++++++++++++++++++++-----
 library/core/src/result.rs       |  2 --
 2 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/library/core/src/alloc/layout.rs b/library/core/src/alloc/layout.rs
index 549a4bc6727fc..c2521ffe9cc96 100644
--- a/library/core/src/alloc/layout.rs
+++ b/library/core/src/alloc/layout.rs
@@ -6,7 +6,7 @@
 
 use crate::error::Error;
 use crate::ptr::{Alignment, NonNull};
-use crate::{cmp, fmt, mem};
+use crate::{assert_unsafe_precondition, cmp, fmt, mem};
 
 // While this function is used in one place and its implementation
 // could be inlined, the previous attempts to do so made rustc
@@ -66,12 +66,25 @@ impl Layout {
     #[inline]
     #[rustc_allow_const_fn_unstable(ptr_alignment_type)]
     pub const fn from_size_align(size: usize, align: usize) -> Result<Self, LayoutError> {
-        if !align.is_power_of_two() {
-            return Err(LayoutError);
+        if Layout::is_size_align_valid(size, align) {
+            // SAFETY: Layout::is_size_align_valid checks the preconditions for this call.
+            let layout = unsafe { Layout::from_size_align_unchecked(size, align) };
+            Ok(layout)
+        } else {
+            Err(LayoutError)
         }
+    }
 
-        // SAFETY: just checked that align is a power of two.
-        Layout::from_size_alignment(size, unsafe { Alignment::new_unchecked(align) })
+    const fn is_size_align_valid(size: usize, align: usize) -> bool {
+        if !align.is_power_of_two() {
+            return false;
+        }
+        // SAFETY: Precondition checked directly above.
+        let align = unsafe { Alignment::new_unchecked(align) };
+        if size > Self::max_size_for_align(align) {
+            return false;
+        }
+        true
     }
 
     #[inline(always)]
@@ -116,6 +129,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::is_size_align_valid(size, align)
+        );
         // SAFETY: the caller is required to uphold the preconditions.
         unsafe { Layout { size, align: Alignment::new_unchecked(align) } }
     }
diff --git a/library/core/src/result.rs b/library/core/src/result.rs
index 7f278296b7b88..73b11f803d929 100644
--- a/library/core/src/result.rs
+++ b/library/core/src/result.rs
@@ -1481,7 +1481,6 @@ impl<T, E> Result<T, E> {
     #[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<T, E> Result<T, E> {
     #[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 e6b0f27e00c0a6e0bb137bab837fbd5e9d2b7629 Mon Sep 17 00:00:00 2001
From: Ben Kimock <kimockb@gmail.com>
Date: Tue, 20 Aug 2024 18:41:07 -0400
Subject: [PATCH 2/2] Try to golf down the amount of code in Layout

---
 library/core/src/alloc/layout.rs | 11 +++--------
 1 file changed, 3 insertions(+), 8 deletions(-)

diff --git a/library/core/src/alloc/layout.rs b/library/core/src/alloc/layout.rs
index c2521ffe9cc96..ad3f9d8087897 100644
--- a/library/core/src/alloc/layout.rs
+++ b/library/core/src/alloc/layout.rs
@@ -68,19 +68,14 @@ impl Layout {
     pub const fn from_size_align(size: usize, align: usize) -> Result<Self, LayoutError> {
         if Layout::is_size_align_valid(size, align) {
             // SAFETY: Layout::is_size_align_valid checks the preconditions for this call.
-            let layout = unsafe { Layout::from_size_align_unchecked(size, align) };
-            Ok(layout)
+            unsafe { Ok(Layout { size, align: mem::transmute(align) }) }
         } else {
             Err(LayoutError)
         }
     }
 
     const fn is_size_align_valid(size: usize, align: usize) -> bool {
-        if !align.is_power_of_two() {
-            return false;
-        }
-        // SAFETY: Precondition checked directly above.
-        let align = unsafe { Alignment::new_unchecked(align) };
+        let Some(align) = Alignment::new(align) else { return false };
         if size > Self::max_size_for_align(align) {
             return false;
         }
@@ -139,7 +134,7 @@ impl Layout {
             ) => Layout::is_size_align_valid(size, align)
         );
         // SAFETY: the caller is required to uphold the preconditions.
-        unsafe { Layout { size, align: Alignment::new_unchecked(align) } }
+        unsafe { Layout { size, align: mem::transmute(align) } }
     }
 
     /// The minimum size in bytes for a memory block of this layout.