Eliminate UbCheck for non-standard libraries

Author: dianqk

compiler/rustc_feature/src/builtin_attrs.rs

@@ -772,6 +772,11 @@ pub const BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
         rustc_allow_incoherent_impl, AttributeType::Normal, template!(Word), ErrorFollowing, @only_local: true,
         "#[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,
+        @only_local: true,
+        "`#![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,

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,18 @@ 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();
+        let debug_assertions = tcx.sess.opts.debug_assertions;
         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, debug_assertions);
+                        }
                         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 +150,19 @@ impl<'tcx> InstSimplifyContext<'tcx, '_> {
         }
     }
 
+    fn simplify_ub_check(
+        &self,
+        source_info: &SourceInfo,
+        rvalue: &mut Rvalue<'tcx>,
+        debug_assertions: bool,
+    ) {
+        if let Rvalue::NullaryOp(NullOp::UbCheck(_), _) = *rvalue {
+            let const_ = Const::from_bool(self.tcx, 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);

compiler/rustc_span/src/symbol.rs

@@ -1560,6 +1560,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,

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)]

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/",

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: UbCheck(LanguageUb)
+    // CHECK: [[assume:_.*]] = const false;
+    // CHECK-NEXT: assume([[assume]]);
+    // CHECK-NEXT: unreachable_unchecked::precondition_check
+    unsafe { x.unwrap_unchecked() }
+}
+
+fn main() {
+    unwrap_unchecked(None);
+}

tests/mir-opt/instsimplify/ub_check.unwrap_unchecked.InstSimplify.diff

@@ -0,0 +1,55 @@
+- // 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 {
+                      }
+                  }
+              }
+          }
+      }
+  
+      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 = UbCheck(LanguageUb);
++         _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;
+      }
+  }
+