Eliminate UbCheck for non-standard libraries

Author: dianqk

Diff for: compiler/rustc_feature/src/builtin_attrs.rs

@@ -821,6 +821,10 @@ pub const BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
         rustc_allow_incoherent_impl, AttributeType::Normal, template!(Word), ErrorFollowing, EncodeCrossCrate::No,
         "#[rustc_allow_incoherent_impl] has to be added to all impl items of an incoherent inherent impl."
     ),
+    rustc_attr!(
+        rustc_preserve_ub_checks, AttributeType::CrateLevel, template!(Word), ErrorFollowing, EncodeCrossCrate::No,
+        "`#![rustc_preserve_ub_checks]` prevents the designated crate from evaluating whether UB checks are enabled when optimizing MIR",
+    ),
     rustc_attr!(
         rustc_deny_explicit_impl,
         AttributeType::Normal,

Diff for: compiler/rustc_middle/src/mir/syntax.rs

@@ -1361,8 +1361,8 @@ pub enum NullOp<'tcx> {
     AlignOf,
     /// Returns the offset of a field
     OffsetOf(&'tcx List<(VariantIdx, FieldIdx)>),
-    /// Returns whether we want to check for UB.
-    /// This returns the value of `cfg!(debug_assertions)` at monomorphization time.
+    /// Returns whether we should perform some UB-checking at runtime.
+    /// See the `ub_checks` intrinsic docs for details.
     UbChecks,
 }
 

Diff for: compiler/rustc_mir_transform/src/instsimplify.rs

@@ -1,10 +1,12 @@
 //! Performs various peephole optimizations.
 
 use crate::simplify::simplify_duplicate_switch_targets;
+use rustc_ast::attr;
 use rustc_middle::mir::*;
 use rustc_middle::ty::layout;
 use rustc_middle::ty::layout::ValidityRequirement;
 use rustc_middle::ty::{self, GenericArgsRef, ParamEnv, Ty, TyCtxt};
+use rustc_span::sym;
 use rustc_span::symbol::Symbol;
 use rustc_target::abi::FieldIdx;
 use rustc_target::spec::abi::Abi;
@@ -22,10 +24,17 @@ impl<'tcx> MirPass<'tcx> for InstSimplify {
             local_decls: &body.local_decls,
             param_env: tcx.param_env_reveal_all_normalized(body.source.def_id()),
         };
+        // FIXME(#116171) Coverage related, also see `unreachable_prop.rs`.
+        let preserve_ub_checks =
+            attr::contains_name(tcx.hir().krate_attrs(), sym::rustc_preserve_ub_checks)
+                || tcx.sess.instrument_coverage();
         for block in body.basic_blocks.as_mut() {
             for statement in block.statements.iter_mut() {
                 match statement.kind {
                     StatementKind::Assign(box (_place, ref mut rvalue)) => {
+                        if !preserve_ub_checks {
+                            ctx.simplify_ub_check(&statement.source_info, rvalue);
+                        }
                         ctx.simplify_bool_cmp(&statement.source_info, rvalue);
                         ctx.simplify_ref_deref(&statement.source_info, rvalue);
                         ctx.simplify_len(&statement.source_info, rvalue);
@@ -140,6 +149,14 @@ impl<'tcx> InstSimplifyContext<'tcx, '_> {
         }
     }
 
+    fn simplify_ub_check(&self, source_info: &SourceInfo, rvalue: &mut Rvalue<'tcx>) {
+        if let Rvalue::NullaryOp(NullOp::UbChecks, _) = *rvalue {
+            let const_ = Const::from_bool(self.tcx, self.tcx.sess.opts.debug_assertions);
+            let constant = ConstOperand { span: source_info.span, const_, user_ty: None };
+            *rvalue = Rvalue::Use(Operand::Constant(Box::new(constant)));
+        }
+    }
+
     fn simplify_cast(&self, rvalue: &mut Rvalue<'tcx>) {
         if let Rvalue::Cast(kind, operand, cast_ty) = rvalue {
             let operand_ty = operand.ty(self.local_decls, self.tcx);

Diff for: compiler/rustc_span/src/symbol.rs

@@ -1572,6 +1572,7 @@ symbols! {
         rustc_peek_maybe_init,
         rustc_peek_maybe_uninit,
         rustc_polymorphize_error,
+        rustc_preserve_ub_checks,
         rustc_private,
         rustc_proc_macro_decls,
         rustc_promotable,

Diff for: library/alloc/src/lib.rs

@@ -176,6 +176,7 @@
 // Language features:
 // tidy-alphabetical-start
 #![cfg_attr(bootstrap, feature(associated_type_bounds))]
+#![cfg_attr(not(bootstrap), rustc_preserve_ub_checks)]
 #![cfg_attr(not(test), feature(coroutine_trait))]
 #![cfg_attr(test, feature(panic_update_hook))]
 #![cfg_attr(test, feature(test))]

Diff for: library/core/src/intrinsics.rs

@@ -2686,12 +2686,14 @@ pub const unsafe fn typed_swap<T>(x: *mut T, y: *mut T) {
     unsafe { ptr::swap_nonoverlapping(x, y, 1) };
 }
 
-/// Returns whether we should perform some UB-checking at runtime. This evaluate to the value of
-/// `cfg!(debug_assertions)` during monomorphization.
-///
-/// This intrinsic is evaluated after monomorphization, which is relevant when mixing crates
-/// compiled with and without debug_assertions. The common case here is a user program built with
-/// debug_assertions linked against the distributed sysroot which is built without debug_assertions.
+/// Returns whether we should perform some UB-checking at runtime. This eventually evaluates to
+/// `cfg!(debug_assertions)`, but behaves different from `cfg!` when mixing crates built with different
+/// flags: if the crate has debug assertions enabled or carries the `#[rustc_preserve_ub_checks]`
+/// attribute, evaluation is delayed until monomorphization (or until the call gets inlined into
+/// a crate that does not delay evaluation further); otherwise it can happen any time.
+///
+/// The common case here is a user program built with debug_assertions linked against the distributed
+/// sysroot which is built without debug_assertions but with `#[rustc_preserve_ub_checks]`.
 /// For code that gets monomorphized in the user crate (i.e., generic functions and functions with
 /// `#[inline]`), gating assertions on `ub_checks()` rather than `cfg!(debug_assertions)` means that
 /// assertions are enabled whenever the *user crate* has debug assertions enabled. However if the

Diff for: library/core/src/lib.rs

@@ -94,6 +94,7 @@
 ))]
 #![no_core]
 #![rustc_coherence_is_core]
+#![cfg_attr(not(bootstrap), rustc_preserve_ub_checks)]
 //
 // Lints:
 #![deny(rust_2021_incompatible_or_patterns)]

Diff for: library/std/src/lib.rs

@@ -221,6 +221,7 @@
 //
 #![cfg_attr(not(feature = "restricted-std"), stable(feature = "rust1", since = "1.0.0"))]
 #![cfg_attr(feature = "restricted-std", unstable(feature = "restricted_std", issue = "none"))]
+#![cfg_attr(not(bootstrap), rustc_preserve_ub_checks)]
 #![doc(
     html_playground_url = "https://play.rust-lang.org/",
     issue_tracker_base_url = "https://github.com/rust-lang/rust/issues/",

Diff for: tests/mir-opt/inline/unwrap_unchecked.unwrap_unchecked.PreCodegen.after.panic-abort.mir

@@ -11,8 +11,6 @@ fn unwrap_unchecked(_1: Option<T>) -> T {
         }
         scope 3 {
             scope 4 (inlined unreachable_unchecked) {
-                let mut _3: bool;
-                let _4: ();
                 scope 5 {
                 }
                 scope 6 (inlined core::ub_checks::check_language_ub) {
@@ -26,23 +24,16 @@ fn unwrap_unchecked(_1: Option<T>) -> T {
     bb0: {
         StorageLive(_2);
         _2 = discriminant(_1);
-        switchInt(move _2) -> [0: bb1, 1: bb2, otherwise: bb3];
+        switchInt(move _2) -> [0: bb2, 1: bb1, otherwise: bb2];
     }
 
     bb1: {
-        StorageLive(_3);
-        _3 = UbChecks();
-        assume(_3);
-        _4 = unreachable_unchecked::precondition_check() -> [return: bb3, unwind unreachable];
-    }
-
-    bb2: {
         _0 = ((_1 as Some).0: T);
         StorageDead(_2);
         return;
     }
 
-    bb3: {
+    bb2: {
         unreachable;
     }
 }

Diff for: tests/mir-opt/inline/unwrap_unchecked.unwrap_unchecked.PreCodegen.after.panic-unwind.mir

@@ -11,8 +11,6 @@ fn unwrap_unchecked(_1: Option<T>) -> T {
         }
         scope 3 {
             scope 4 (inlined unreachable_unchecked) {
-                let mut _3: bool;
-                let _4: ();
                 scope 5 {
                 }
                 scope 6 (inlined core::ub_checks::check_language_ub) {
@@ -26,23 +24,16 @@ fn unwrap_unchecked(_1: Option<T>) -> T {
     bb0: {
         StorageLive(_2);
         _2 = discriminant(_1);
-        switchInt(move _2) -> [0: bb1, 1: bb2, otherwise: bb3];
+        switchInt(move _2) -> [0: bb2, 1: bb1, otherwise: bb2];
     }
 
     bb1: {
-        StorageLive(_3);
-        _3 = UbChecks();
-        assume(_3);
-        _4 = unreachable_unchecked::precondition_check() -> [return: bb3, unwind unreachable];
-    }
-
-    bb2: {
         _0 = ((_1 as Some).0: T);
         StorageDead(_2);
         return;
     }
 
-    bb3: {
+    bb2: {
         unreachable;
     }
 }

Diff for: tests/mir-opt/instsimplify/ub_check.rs

@@ -0,0 +1,16 @@
+//@ unit-test: InstSimplify
+//@ compile-flags: -Cdebug-assertions=no -Zinline-mir
+
+// EMIT_MIR ub_check.unwrap_unchecked.InstSimplify.diff
+pub fn unwrap_unchecked(x: Option<i32>) -> i32 {
+    // CHECK-LABEL: fn unwrap_unchecked(
+    // CHECK-NOT: UbChecks()
+    // CHECK: [[assume:_.*]] = const false;
+    // CHECK-NEXT: assume([[assume]]);
+    // CHECK-NEXT: unreachable_unchecked::precondition_check
+    unsafe { x.unwrap_unchecked() }
+}
+
+fn main() {
+    unwrap_unchecked(None);
+}

Diff for: tests/mir-opt/instsimplify/ub_check.unwrap_unchecked.InstSimplify.diff

@@ -0,0 +1,59 @@
+- // MIR for `unwrap_unchecked` before InstSimplify
++ // MIR for `unwrap_unchecked` after InstSimplify
+  
+  fn unwrap_unchecked(_1: Option<i32>) -> i32 {
+      debug x => _1;
+      let mut _0: i32;
+      let mut _2: std::option::Option<i32>;
+      scope 1 {
+          scope 2 (inlined #[track_caller] Option::<i32>::unwrap_unchecked) {
+              debug self => _2;
+              let mut _3: isize;
+              scope 3 {
+                  debug val => _0;
+              }
+              scope 4 {
+                  scope 5 (inlined unreachable_unchecked) {
+                      let mut _4: bool;
+                      let _5: ();
+                      scope 6 {
+                      }
+                      scope 7 (inlined core::ub_checks::check_language_ub) {
+                          scope 8 (inlined core::ub_checks::check_language_ub::runtime) {
+                          }
+                      }
+                  }
+              }
+          }
+      }
+  
+      bb0: {
+          StorageLive(_2);
+          _2 = _1;
+          StorageLive(_3);
+          StorageLive(_5);
+          _3 = discriminant(_2);
+          switchInt(move _3) -> [0: bb2, 1: bb3, otherwise: bb1];
+      }
+  
+      bb1: {
+          unreachable;
+      }
+  
+      bb2: {
+          StorageLive(_4);
+-         _4 = UbChecks();
++         _4 = const false;
+          assume(_4);
+          _5 = unreachable_unchecked::precondition_check() -> [return: bb1, unwind unreachable];
+      }
+  
+      bb3: {
+          _0 = move ((_2 as Some).0: i32);
+          StorageDead(_5);
+          StorageDead(_3);
+          StorageDead(_2);
+          return;
+      }
+  }
+  

Diff for: tests/mir-opt/pre-codegen/duplicate_switch_targets.ub_if_b.PreCodegen.after.mir

@@ -5,8 +5,6 @@ fn ub_if_b(_1: Thing) -> Thing {
     let mut _0: Thing;
     let mut _2: isize;
     scope 1 (inlined unreachable_unchecked) {
-        let mut _3: bool;
-        let _4: ();
         scope 2 {
         }
         scope 3 (inlined core::ub_checks::check_language_ub) {
@@ -17,7 +15,7 @@ fn ub_if_b(_1: Thing) -> Thing {
 
     bb0: {
         _2 = discriminant(_1);
-        switchInt(move _2) -> [0: bb1, 1: bb2, otherwise: bb3];
+        switchInt(move _2) -> [0: bb1, 1: bb2, otherwise: bb2];
     }
 
     bb1: {
@@ -26,13 +24,6 @@ fn ub_if_b(_1: Thing) -> Thing {
     }
 
     bb2: {
-        StorageLive(_3);
-        _3 = UbChecks();
-        assume(_3);
-        _4 = unreachable_unchecked::precondition_check() -> [return: bb3, unwind unreachable];
-    }
-
-    bb3: {
         unreachable;
     }
 }