From 3389c502b9e6be780892732a29480a21cd1fd16e Mon Sep 17 00:00:00 2001
From: Nilstrieb <48135649+Nilstrieb@users.noreply.github.com>
Date: Wed, 14 Feb 2024 22:08:03 +0100
Subject: [PATCH] Add `#[rustc_no_mir_inline]` for standard library UB checks

Co-authored-by: Ben Kimock <kimockb@gmail.com>
---
 compiler/rustc_feature/src/builtin_attrs.rs |  4 ++++
 compiler/rustc_mir_transform/src/inline.rs  |  4 ++++
 compiler/rustc_span/src/symbol.rs           |  1 +
 library/core/src/intrinsics.rs              | 18 +++++++++++++++---
 4 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/compiler/rustc_feature/src/builtin_attrs.rs b/compiler/rustc_feature/src/builtin_attrs.rs
index 99875ec540545..003baf071b8c8 100644
--- a/compiler/rustc_feature/src/builtin_attrs.rs
+++ b/compiler/rustc_feature/src/builtin_attrs.rs
@@ -792,6 +792,10 @@ pub const BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
         rustc_intrinsic, Normal, template!(Word), ErrorFollowing,
         "the `#[rustc_intrinsic]` attribute is used to declare intrinsics with function bodies",
     ),
+    rustc_attr!(
+        rustc_no_mir_inline, Normal, template!(Word), WarnFollowing,
+        "#[rustc_no_mir_inline] prevents the MIR inliner from inlining a function while not affecting codegen"
+    ),
 
     // ==========================================================================
     // Internal attributes, Testing:
diff --git a/compiler/rustc_mir_transform/src/inline.rs b/compiler/rustc_mir_transform/src/inline.rs
index 2009539d4d084..36546a03cdfc5 100644
--- a/compiler/rustc_mir_transform/src/inline.rs
+++ b/compiler/rustc_mir_transform/src/inline.rs
@@ -421,6 +421,10 @@ impl<'tcx> Inliner<'tcx> {
         callee_attrs: &CodegenFnAttrs,
         cross_crate_inlinable: bool,
     ) -> Result<(), &'static str> {
+        if self.tcx.has_attr(callsite.callee.def_id(), sym::rustc_no_mir_inline) {
+            return Err("#[rustc_no_mir_inline]");
+        }
+
         if let InlineAttr::Never = callee_attrs.inline {
             return Err("never inline hint");
         }
diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs
index 46472a131ff4b..609ab054da21e 100644
--- a/compiler/rustc_span/src/symbol.rs
+++ b/compiler/rustc_span/src/symbol.rs
@@ -1441,6 +1441,7 @@ symbols! {
         rustc_mir,
         rustc_must_implement_one_of,
         rustc_never_returns_null_ptr,
+        rustc_no_mir_inline,
         rustc_nonnull_optimization_guaranteed,
         rustc_nounwind,
         rustc_object_lifetime_default,
diff --git a/library/core/src/intrinsics.rs b/library/core/src/intrinsics.rs
index e42b2315ea5d7..f9d89795a9988 100644
--- a/library/core/src/intrinsics.rs
+++ b/library/core/src/intrinsics.rs
@@ -2706,13 +2706,25 @@ pub const unsafe fn const_deallocate(_ptr: *mut u8, _size: usize, _align: usize)
 macro_rules! assert_unsafe_precondition {
     ($message:expr, ($($name:ident:$ty:ty = $arg:expr),*$(,)?) => $e:expr $(,)?) => {
         {
+            // #[cfg(bootstrap)] (this comment)
             // When the standard library is compiled with debug assertions, we want the check to inline for better performance.
             // This is important when working on the compiler, which is compiled with debug assertions locally.
             // When not compiled with debug assertions (so the precompiled std) we outline the check to minimize the compile
             // time impact when debug assertions are disabled.
-            // It is not clear whether that is the best solution, see #120848.
-            #[cfg_attr(debug_assertions, inline(always))]
-            #[cfg_attr(not(debug_assertions), inline(never))]
+            // The proper solution to this is the `#[rustc_no_mir_inline]` below, but we still want decent performance for cfg(bootstrap).
+            #[cfg_attr(all(debug_assertions, bootstrap), inline(always))]
+            #[cfg_attr(all(not(debug_assertions), bootstrap), inline(never))]
+
+            // This check is inlineable, but not by the MIR inliner.
+            // The reason for this is that the MIR inliner is in an exceptionally bad position
+            // to think about whether or not to inline this. In MIR, this call is gated behind `debug_assertions`,
+            // which will codegen to `false` in release builds. Inlining the check would be wasted work in that case and
+            // would be bad for compile times.
+            //
+            // LLVM on the other hand sees the constant branch, so if it's `false`, it can immediately delete it without
+            // inlining the check. If it's `true`, it can inline it and get significantly better performance.
+            #[cfg_attr(not(bootstrap), rustc_no_mir_inline)]
+            #[cfg_attr(not(bootstrap), inline)]
             #[rustc_nounwind]
             fn precondition_check($($name:$ty),*) {
                 if !$e {