Try enabling precondition checks on ptr::{read,write}

Author: saethlin

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

@@ -1012,6 +1012,10 @@ pub const BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
         rustc_no_mir_inline, Normal, template!(Word), WarnFollowing, EncodeCrossCrate::Yes,
         "#[rustc_no_mir_inline] prevents the MIR inliner from inlining a function while not affecting codegen"
     ),
+    rustc_attr!(
+        rustc_no_ubchecks, Normal, template!(Word), WarnFollowing, EncodeCrossCrate::Yes,
+        "#[rustc_no_ubchecks] asks the compiler to delete UB checks from a function"
+    ),
     rustc_attr!(
         rustc_force_inline, Normal, template!(Word, NameValueStr: "reason"), WarnFollowing, EncodeCrossCrate::Yes,
         "#![rustc_force_inline] forces a free function to be inlined"

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

@@ -30,19 +30,27 @@ impl<'tcx> crate::MirPass<'tcx> for InstSimplify {
     }
 
     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        let def_id = body.source.def_id();
         let ctx = InstSimplifyContext {
             tcx,
             local_decls: &body.local_decls,
             typing_env: body.typing_env(tcx),
         };
         let preserve_ub_checks =
             attr::contains_name(tcx.hir().krate_attrs(), sym::rustc_preserve_ub_checks);
+        let remove_ub_checks = if tcx.is_coroutine(def_id) {
+            false
+        } else {
+            tcx.has_attr(def_id, sym::rustc_no_ubchecks)
+        };
         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(rvalue);
+                        if remove_ub_checks {
+                            ctx.simplify_ub_check(rvalue, false);
+                        } else if !preserve_ub_checks {
+                            ctx.simplify_ub_check(rvalue, tcx.sess.ub_checks());
                         }
                         ctx.simplify_bool_cmp(rvalue);
                         ctx.simplify_ref_deref(rvalue);
@@ -194,9 +202,9 @@ impl<'tcx> InstSimplifyContext<'_, 'tcx> {
         }
     }
 
-    fn simplify_ub_check(&self, rvalue: &mut Rvalue<'tcx>) {
+    fn simplify_ub_check(&self, rvalue: &mut Rvalue<'tcx>, ub_checks: bool) {
         if let Rvalue::NullaryOp(NullOp::UbChecks, _) = *rvalue {
-            let const_ = Const::from_bool(self.tcx, self.tcx.sess.ub_checks());
+            let const_ = Const::from_bool(self.tcx, ub_checks);
             let constant = ConstOperand { span: DUMMY_SP, const_, user_ty: None };
             *rvalue = Rvalue::Use(Operand::Constant(Box::new(constant)));
         }

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

@@ -1758,6 +1758,7 @@ symbols! {
         rustc_never_returns_null_ptr,
         rustc_never_type_options,
         rustc_no_mir_inline,
+        rustc_no_ubchecks,
         rustc_nonnull_optimization_guaranteed,
         rustc_nounwind,
         rustc_object_lifetime_default,

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

@@ -1430,6 +1430,7 @@ impl<T: ?Sized, A: Allocator> Box<T, A> {
     #[must_use = "losing the pointer will leak memory"]
     #[unstable(feature = "allocator_api", issue = "32838")]
     #[inline]
+    #[cfg_attr(not(bootstrap), rustc_no_ubchecks)]
     pub fn into_raw_with_allocator(b: Self) -> (*mut T, A) {
         let mut b = mem::ManuallyDrop::new(b);
         // We carefully get the raw pointer out in a way that Miri's aliasing model understands what

Diff for: library/core/src/mem/manually_drop.rs

@@ -217,6 +217,7 @@ impl<T> ManuallyDrop<T> {
     #[must_use = "if you don't need the value, you can use `ManuallyDrop::drop` instead"]
     #[stable(feature = "manually_drop_take", since = "1.42.0")]
     #[inline]
+    #[cfg_attr(not(bootstrap), rustc_no_ubchecks)]
     pub unsafe fn take(slot: &mut ManuallyDrop<T>) -> T {
         // SAFETY: we are reading from a reference, which is guaranteed
         // to be valid for reads.

Diff for: library/core/src/mem/mod.rs

@@ -857,6 +857,7 @@ pub fn take<T: Default>(dest: &mut T) -> T {
 #[must_use = "if you don't need the old value, you can just assign the new value directly"]
 #[rustc_const_stable(feature = "const_replace", since = "1.83.0")]
 #[cfg_attr(not(test), rustc_diagnostic_item = "mem_replace")]
+#[cfg_attr(not(bootstrap), rustc_no_ubchecks)]
 pub const fn replace<T>(dest: &mut T, src: T) -> T {
     // It may be tempting to use `swap` to avoid `unsafe` here. Don't!
     // The compiler optimizes the implementation below to two `memcpy`s

Diff for: library/core/src/ptr/mod.rs

@@ -1136,6 +1136,7 @@ pub const unsafe fn swap_nonoverlapping<T>(x: *mut T, y: *mut T, count: usize) {
 /// LLVM can vectorize this (at least it can for the power-of-two-sized types
 /// `swap_nonoverlapping` tries to use) so no need to manually SIMD it.
 #[inline]
+#[cfg_attr(not(bootstrap), rustc_no_ubchecks)]
 const unsafe fn swap_nonoverlapping_simple_untyped<T>(x: *mut T, y: *mut T, count: usize) {
     let x = x.cast::<MaybeUninit<T>>();
     let y = y.cast::<MaybeUninit<T>>();
@@ -1367,7 +1368,6 @@ pub const unsafe fn read<T>(src: *const T) -> T {
 
     // SAFETY: the caller must guarantee that `src` is valid for reads.
     unsafe {
-        #[cfg(debug_assertions)] // Too expensive to always enable (for now?)
         ub_checks::assert_unsafe_precondition!(
             check_language_ub,
             "ptr::read requires that the pointer argument is aligned and non-null",
@@ -1569,7 +1569,6 @@ pub const unsafe fn write<T>(dst: *mut T, src: T) {
     // `dst` cannot overlap `src` because the caller has mutable access
     // to `dst` while `src` is owned by this function.
     unsafe {
-        #[cfg(debug_assertions)] // Too expensive to always enable (for now?)
         ub_checks::assert_unsafe_precondition!(
             check_language_ub,
             "ptr::write requires that the pointer argument is aligned and non-null",

Diff for: tests/codegen/mem-replace-big-type.rs

@@ -4,8 +4,6 @@
 // known to be `1` after inlining).
 
 //@ compile-flags: -C no-prepopulate-passes -Zinline-mir=no
-//@ ignore-std-debug-assertions
-// Reason: precondition checks in ptr::read make them a bad candidate for MIR inlining
 //@ needs-deterministic-layouts
 
 #![crate_type = "lib"]

Diff for: tests/codegen/mem-replace-simple-type.rs

@@ -1,7 +1,5 @@
 //@ compile-flags: -O -C no-prepopulate-passes
 //@ only-x86_64 (to not worry about usize differing)
-//@ ignore-std-debug-assertions
-// Reason: precondition checks make mem::replace not a candidate for MIR inlining
 
 #![crate_type = "lib"]
 

Diff for: tests/mir-opt/pre-codegen/mem_replace.mem_replace.PreCodegen.after.panic-abort.mir

@@ -6,10 +6,22 @@ fn mem_replace(_1: &mut u32, _2: u32) -> u32 {
     let mut _0: u32;
     scope 1 (inlined std::mem::replace::<u32>) {
         scope 2 {
-            scope 4 (inlined std::ptr::write::<u32>) {
+            scope 7 (inlined std::ptr::write::<u32>) {
+                scope 8 (inlined core::ub_checks::check_language_ub) {
+                    scope 9 (inlined core::ub_checks::check_language_ub::runtime) {
+                    }
+                }
+                scope 10 (inlined align_of::<u32>) {
+                }
             }
         }
         scope 3 (inlined std::ptr::read::<u32>) {
+            scope 4 (inlined core::ub_checks::check_language_ub) {
+                scope 5 (inlined core::ub_checks::check_language_ub::runtime) {
+                }
+            }
+            scope 6 (inlined align_of::<u32>) {
+            }
         }
     }
 

Diff for: tests/mir-opt/pre-codegen/mem_replace.mem_replace.PreCodegen.after.panic-unwind.mir

@@ -6,10 +6,22 @@ fn mem_replace(_1: &mut u32, _2: u32) -> u32 {
     let mut _0: u32;
     scope 1 (inlined std::mem::replace::<u32>) {
         scope 2 {
-            scope 4 (inlined std::ptr::write::<u32>) {
+            scope 7 (inlined std::ptr::write::<u32>) {
+                scope 8 (inlined core::ub_checks::check_language_ub) {
+                    scope 9 (inlined core::ub_checks::check_language_ub::runtime) {
+                    }
+                }
+                scope 10 (inlined align_of::<u32>) {
+                }
             }
         }
         scope 3 (inlined std::ptr::read::<u32>) {
+            scope 4 (inlined core::ub_checks::check_language_ub) {
+                scope 5 (inlined core::ub_checks::check_language_ub::runtime) {
+                }
+            }
+            scope 6 (inlined align_of::<u32>) {
+            }
         }
     }
 

Diff for: tests/mir-opt/pre-codegen/mem_replace.rs

@@ -1,7 +1,5 @@
 // skip-filecheck
 //@ compile-flags: -O -C debuginfo=0 -Zmir-opt-level=2 -Zinline-mir
-//@ ignore-std-debug-assertions
-// Reason: precondition checks on ptr::read/write are under cfg(debug_assertions)
 // EMIT_MIR_FOR_EACH_PANIC_STRATEGY
 
 #![crate_type = "lib"]

Diff for: tests/mir-opt/pre-codegen/ptr_offset.rs

@@ -1,6 +1,5 @@
 // skip-filecheck
 //@ compile-flags: -O -C debuginfo=0 -Zmir-opt-level=2 -Zinline-mir
-//@ ignore-std-debug-assertions (precondition checks are under cfg(debug_assertions))
 // EMIT_MIR_FOR_EACH_PANIC_STRATEGY
 
 #![crate_type = "lib"]