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

Author: saethlin

compiler/rustc_feature/src/builtin_attrs.rs

@@ -1408,6 +1408,10 @@ pub static 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"

compiler/rustc_mir_transform/src/instsimplify.rs

@@ -29,21 +29,29 @@ 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() {
                 let StatementKind::Assign(box (.., rvalue)) = &mut statement.kind else {
                     continue;
                 };
 
-                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);
@@ -168,13 +176,13 @@ 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) {
         // FIXME: Should we do the same for overflow checks?
         let Rvalue::NullaryOp(NullOp::RuntimeChecks(RuntimeChecks::UbChecks)) = *rvalue else {
             return;
         };
 
-        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)));
     }

compiler/rustc_span/src/symbol.rs

@@ -1961,6 +1961,7 @@ symbols! {
         rustc_no_implicit_autorefs,
         rustc_no_implicit_bounds,
         rustc_no_mir_inline,
+        rustc_no_ubchecks,
         rustc_nonnull_optimization_guaranteed,
         rustc_nounwind,
         rustc_objc_class,

library/alloc/src/boxed.rs

@@ -1475,6 +1475,7 @@ impl<T: ?Sized, A: Allocator> Box<T, A> {
     #[must_use = "losing the pointer will leak memory"]
     #[unstable(feature = "allocator_api", issue = "32838")]
     #[inline]
+    #[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

library/core/src/mem/manually_drop.rs

@@ -219,6 +219,7 @@ impl<T> ManuallyDrop<T> {
     #[stable(feature = "manually_drop_take", since = "1.42.0")]
     #[rustc_const_unstable(feature = "const_manually_drop_take", issue = "148773")]
     #[inline]
+    #[rustc_no_ubchecks]
     pub const unsafe fn take(slot: &mut ManuallyDrop<T>) -> T {
         // SAFETY: we are reading from a reference, which is guaranteed
         // to be valid for reads.

library/core/src/mem/mod.rs

@@ -876,6 +876,7 @@ pub const fn take<T: [const] 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")]
 #[rustc_diagnostic_item = "mem_replace"]
+#[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

library/core/src/ptr/mod.rs

@@ -1400,6 +1400,7 @@ pub const unsafe fn swap_nonoverlapping<T>(x: *mut T, y: *mut T, count: usize) {
 
 /// Same behavior and safety conditions as [`swap_nonoverlapping`]
 #[inline]
+#[rustc_no_ubchecks]
 const unsafe fn swap_nonoverlapping_const<T>(x: *mut T, y: *mut T, count: usize) {
     let mut i = 0;
     while i < count {
@@ -1695,7 +1696,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",
@@ -1895,7 +1895,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",

tests/codegen-llvm/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"]

tests/codegen-llvm/mem-replace-simple-type.rs

@@ -1,7 +1,5 @@
 //@ compile-flags: -Copt-level=3 -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"]
 

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