Eliminate UbCheck for non-standard libraries

Author: dianqk

compiler/rustc_feature/src/builtin_attrs.rs

@@ -468,6 +468,11 @@ pub const BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
         no_core, CrateLevel, template!(Word), WarnFollowing,
         @only_local: true, experimental!(no_core)
     ),
+    gated!(
+        rustc_ub_check, CrateLevel, template!(Word), WarnFollowing,
+        @only_local: true,
+        "`#![rustc_ub_check]` is an internal feature of the standard library, mark don't optimize UB Checks",
+    ),
     // RFC 2412
     gated!(
         optimize, Normal, template!(List: "size|speed"), ErrorPreceding,

compiler/rustc_feature/src/unstable.rs

@@ -221,6 +221,8 @@ declare_features! (
     (internal, profiler_runtime, "1.18.0", None),
     /// Allows using `rustc_*` attributes (RFC 572).
     (internal, rustc_attrs, "1.0.0", None),
+    /// Keep UB checks for the standard library.
+    (internal, rustc_ub_check, "CURRENT_RUSTC_VERSION", None),
     /// Allows using the `#[stable]` and `#[unstable]` attributes.
     (internal, staged_api, "1.0.0", None),
     /// Added for testing unstable lints; perma-unstable.

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,15 @@ impl<'tcx> MirPass<'tcx> for InstSimplify {
             local_decls: &body.local_decls,
             param_env: tcx.param_env_reveal_all_normalized(body.source.def_id()),
         };
+        let rustc_ub_check = attr::contains_name(tcx.hir().krate_attrs(), sym::rustc_ub_check);
+        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 !rustc_ub_check {
+                            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 +147,24 @@ impl<'tcx> InstSimplifyContext<'tcx, '_> {
         }
     }
 
+    fn simplify_ub_check(
+        &self,
+        source_info: &SourceInfo,
+        rvalue: &mut Rvalue<'tcx>,
+        debug_assertions: bool,
+    ) {
+        if let Rvalue::NullaryOp(ref null_op, _) = *rvalue {
+            match null_op {
+                NullOp::SizeOf | NullOp::AlignOf | NullOp::OffsetOf(_) => {}
+                NullOp::UbCheck(_) => {
+                    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

@@ -1576,6 +1576,7 @@ symbols! {
         rustc_test_marker,
         rustc_then_this_would_need,
         rustc_trivial_field_reads,
+        rustc_ub_check,
         rustc_unsafe_specialization_marker,
         rustc_variance,
         rustc_variance_of_opaques,

library/core/src/lib.rs

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

library/std/src/lib.rs

@@ -238,6 +238,8 @@
 #![no_std]
 // Tell the compiler to link to either panic_abort or panic_unwind
 #![needs_panic_runtime]
+#![cfg_attr(not(bootstrap), feature(rustc_ub_check))]
+#![cfg_attr(not(bootstrap), rustc_ub_check)]
 //
 // Lints:
 #![warn(deprecated_in_future)]

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;
+      }
+  }
+  

tests/ui/feature-gates/feature-gate-rustc-ub-check.rs

@@ -0,0 +1,2 @@
+#![rustc_ub_check] //~ERROR `#![rustc_ub_check]` is an internal feature of the standard library, mark don't optimize UB Checks
+fn main() {}

tests/ui/feature-gates/feature-gate-rustc-ub-check.stderr

@@ -0,0 +1,12 @@
+error[E0658]: `#![rustc_ub_check]` is an internal feature of the standard library, mark don't optimize UB Checks
+  --> $DIR/feature-gate-rustc-ub-check.rs:1:1
+   |
+LL | #![rustc_ub_check]
+   | ^^^^^^^^^^^^^^^^^^
+   |
+   = help: add `#![feature(rustc_ub_check)]` to the crate attributes to enable
+   = note: this compiler was built on YYYY-MM-DD; consider upgrading it if it is out of date
+
+error: aborting due to 1 previous error
+
+For more information about this error, try `rustc --explain E0658`.