Auto merge of #111010 - scottmcm:mem-replace-simpler, r=WaffleLapkin

Make `mem::replace` simpler in codegen

Since they'd mentioned more intrinsics for simplifying stuff recently,
r? `@WaffleLapkin`

This is a continuation of me looking at foundational stuff that ends up with more instructions than it really needs.  Specifically I noticed this one because `Range::next` isn't MIR-inlining, and one of the largest parts of it is a `replace::<usize>` that's a good dozen instructions instead of the two it could be.

So this means that `ptr::write` with a `Copy` type no longer generates worse IR than manually dereferencing (well, at least in LLVM -- MIR still has bonus pointer casts), and in doing so means that we're finally down to just the two essential `memcpy`s when emitting `mem::replace` for a large type, rather than the bonus-`alloca` and three `memcpy`s we emitted before this ([or the 6 we currently emit in 1.69 stable](https://rust.godbolt.org/z/67W8on6nP)).  That said, LLVM does _usually_ manage to optimize the extra code away.  But it's still nice for it not to have to do as much, thanks to (for example) not going through an `alloca` when `replace`ing a primitive like a `usize`.

(This is a new intrinsic, but one that's immediately lowered to existing MIR constructs, so not anything that MIRI or the codegen backends or MIR semantics needs to do work to handle.)

Author: bors

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -381,6 +381,7 @@ pub fn check_intrinsic_type(tcx: TyCtxt<'_>, it: &hir::ForeignItem<'_>) {
             sym::unlikely => (0, vec![tcx.types.bool], tcx.types.bool),
 
             sym::read_via_copy => (1, vec![tcx.mk_imm_ptr(param(0))], param(0)),
+            sym::write_via_move => (1, vec![tcx.mk_mut_ptr(param(0)), param(0)], tcx.mk_unit()),
 
             sym::discriminant_value => {
                 let assoc_items = tcx.associated_item_def_ids(

compiler/rustc_mir_transform/src/lower_intrinsics.rs

@@ -179,6 +179,29 @@ impl<'tcx> MirPass<'tcx> for LowerIntrinsics {
                             }
                         }
                     }
+                    sym::write_via_move => {
+                        let target = target.unwrap();
+                        let Ok([ptr, val]) = <[_; 2]>::try_from(std::mem::take(args)) else {
+                            span_bug!(
+                                terminator.source_info.span,
+                                "Wrong number of arguments for write_via_move intrinsic",
+                            );
+                        };
+                        let derefed_place =
+                            if let Some(place) = ptr.place() && let Some(local) = place.as_local() {
+                                tcx.mk_place_deref(local.into())
+                            } else {
+                                span_bug!(terminator.source_info.span, "Only passing a local is supported");
+                            };
+                        block.statements.push(Statement {
+                            source_info: terminator.source_info,
+                            kind: StatementKind::Assign(Box::new((
+                                derefed_place,
+                                Rvalue::Use(val),
+                            ))),
+                        });
+                        terminator.kind = TerminatorKind::Goto { target };
+                    }
                     sym::discriminant_value => {
                         if let (Some(target), Some(arg)) = (*target, args[0].place()) {
                             let arg = tcx.mk_place_deref(arg);

compiler/rustc_span/src/symbol.rs

@@ -1636,6 +1636,7 @@ symbols! {
         write_bytes,
         write_macro,
         write_str,
+        write_via_move,
         writeln_macro,
         x87_reg,
         xer,

library/core/src/intrinsics.rs

@@ -2257,12 +2257,23 @@ extern "rust-intrinsic" {
     /// This is an implementation detail of [`crate::ptr::read`] and should
     /// not be used anywhere else.  See its comments for why this exists.
     ///
-    /// This intrinsic can *only* be called where the argument is a local without
-    /// projections (`read_via_copy(p)`, not `read_via_copy(*p)`) so that it
+    /// This intrinsic can *only* be called where the pointer is a local without
+    /// projections (`read_via_copy(ptr)`, not `read_via_copy(*ptr)`) so that it
     /// trivially obeys runtime-MIR rules about derefs in operands.
     #[rustc_const_unstable(feature = "const_ptr_read", issue = "80377")]
     #[rustc_nounwind]
-    pub fn read_via_copy<T>(p: *const T) -> T;
+    pub fn read_via_copy<T>(ptr: *const T) -> T;
+
+    /// This is an implementation detail of [`crate::ptr::write`] and should
+    /// not be used anywhere else.  See its comments for why this exists.
+    ///
+    /// This intrinsic can *only* be called where the pointer is a local without
+    /// projections (`write_via_move(ptr, x)`, not `write_via_move(*ptr, x)`) so
+    /// that it trivially obeys runtime-MIR rules about derefs in operands.
+    #[cfg(not(bootstrap))]
+    #[rustc_const_unstable(feature = "const_ptr_write", issue = "86302")]
+    #[rustc_nounwind]
+    pub fn write_via_move<T>(ptr: *mut T, value: T);
 
     /// Returns the value of the discriminant for the variant in 'v';
     /// if `T` has no discriminant, returns `0`.
@@ -2828,3 +2839,16 @@ pub const unsafe fn transmute_unchecked<Src, Dst>(src: Src) -> Dst {
     // SAFETY: It's a transmute -- the caller promised it's fine.
     unsafe { transmute_copy(&ManuallyDrop::new(src)) }
 }
+
+/// Polyfill for bootstrap
+#[cfg(bootstrap)]
+pub const unsafe fn write_via_move<T>(ptr: *mut T, value: T) {
+    use crate::mem::*;
+    // SAFETY: the caller must guarantee that `dst` is valid for writes.
+    // `dst` cannot overlap `src` because the caller has mutable access
+    // to `dst` while `src` is owned by this function.
+    unsafe {
+        copy_nonoverlapping::<T>(&value, ptr, 1);
+        forget(value);
+    }
+}

library/core/src/ptr/mod.rs

@@ -1349,13 +1349,13 @@ pub const unsafe fn read_unaligned<T>(src: *const T) -> T {
 #[rustc_const_unstable(feature = "const_ptr_write", issue = "86302")]
 #[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
 pub const unsafe fn write<T>(dst: *mut T, src: T) {
-    // We are calling the intrinsics directly to avoid function calls in the generated code
-    // as `intrinsics::copy_nonoverlapping` is a wrapper function.
-    extern "rust-intrinsic" {
-        #[rustc_const_stable(feature = "const_intrinsic_copy", since = "1.63.0")]
-        #[rustc_nounwind]
-        fn copy_nonoverlapping<T>(src: *const T, dst: *mut T, count: usize);
-    }
+    // Semantically, it would be fine for this to be implemented as a
+    // `copy_nonoverlapping` and appropriate drop suppression of `src`.
+
+    // However, implementing via that currently produces more MIR than is ideal.
+    // Using an intrinsic keeps it down to just the simple `*dst = move src` in
+    // MIR (11 statements shorter, at the time of writing), and also allows
+    // `src` to stay an SSA value in codegen_ssa, rather than a memory one.
 
     // SAFETY: the caller must guarantee that `dst` is valid for writes.
     // `dst` cannot overlap `src` because the caller has mutable access
@@ -1365,8 +1365,7 @@ pub const unsafe fn write<T>(dst: *mut T, src: T) {
             "ptr::write requires that the pointer argument is aligned and non-null",
             [T](dst: *mut T) => is_aligned_and_not_null(dst)
         );
-        copy_nonoverlapping(&src as *const T, dst, 1);
-        intrinsics::forget(src);
+        intrinsics::write_via_move(dst, src)
     }
 }
 

src/tools/miri/tests/fail/dangling_pointers/null_pointer_write_zst.rs

@@ -7,5 +7,5 @@ fn main() {
     // Also not assigning directly as that's array initialization, not assignment.
     let zst_val = [1u8; 0];
     unsafe { std::ptr::null_mut::<[u8; 0]>().write(zst_val) };
-    //~^ERROR: memory access failed: null pointer is a dangling pointer
+    //~^ERROR: dereferencing pointer failed: null pointer is a dangling pointer
 }

src/tools/miri/tests/fail/dangling_pointers/null_pointer_write_zst.stderr

@@ -1,8 +1,8 @@
-error: Undefined Behavior: memory access failed: null pointer is a dangling pointer (it has no provenance)
+error: Undefined Behavior: dereferencing pointer failed: null pointer is a dangling pointer (it has no provenance)
   --> $DIR/null_pointer_write_zst.rs:LL:CC
    |
 LL |     unsafe { std::ptr::null_mut::<[u8; 0]>().write(zst_val) };
-   |              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ memory access failed: null pointer is a dangling pointer (it has no provenance)
+   |              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ dereferencing pointer failed: null pointer is a dangling pointer (it has no provenance)
    |
    = help: this indicates a bug in the program: it performed an invalid operation, and caused Undefined Behavior
    = help: see https://doc.rust-lang.org/nightly/reference/behavior-considered-undefined.html for further information

tests/codegen/mem-replace-big-type.rs

@@ -11,7 +11,9 @@
 #[repr(C, align(8))]
 pub struct Big([u64; 7]);
 pub fn replace_big(dst: &mut Big, src: Big) -> Big {
-    // Before the `read_via_copy` intrinsic, this emitted six `memcpy`s.
+    // Back in 1.68, this emitted six `memcpy`s.
+    // `read_via_copy` in 1.69 got that down to three.
+    // `write_via_move` has it down to just the two essential ones.
     std::mem::replace(dst, src)
 }
 
@@ -20,17 +22,13 @@ pub fn replace_big(dst: &mut Big, src: Big) -> Big {
 
 // CHECK-NOT: call void @llvm.memcpy
 
-// For a large type, we expect exactly three `memcpy`s
+// For a large type, we expect exactly two `memcpy`s
 // CHECK-LABEL: define internal void @{{.+}}mem{{.+}}replace{{.+}}sret(%Big)
     // CHECK-NOT: alloca
-    // CHECK: alloca %Big
-    // CHECK-NOT: alloca
-    // CHECK-NOT: call void @llvm.memcpy
-    // CHECK: call void @llvm.memcpy.{{.+}}({{i8\*|ptr}} align 8 %{{.*}}, {{i8\*|ptr}} align 8 %{{.*}}, i{{.*}} 56, i1 false)
     // CHECK-NOT: call void @llvm.memcpy
-    // CHECK: call void @llvm.memcpy.{{.+}}({{i8\*|ptr}} align 8 %{{.*}}, {{i8\*|ptr}} align 8 %{{.*}}, i{{.*}} 56, i1 false)
+    // CHECK: call void @llvm.memcpy.{{.+}}({{i8\*|ptr}} align 8 %0, {{i8\*|ptr}} align 8 %dest, i{{.*}} 56, i1 false)
     // CHECK-NOT: call void @llvm.memcpy
-    // CHECK: call void @llvm.memcpy.{{.+}}({{i8\*|ptr}} align 8 %{{.*}}, {{i8\*|ptr}} align 8 %{{.*}}, i{{.*}} 56, i1 false)
+    // CHECK: call void @llvm.memcpy.{{.+}}({{i8\*|ptr}} align 8 %dest, {{i8\*|ptr}} align 8 %src, i{{.*}} 56, i1 false)
     // CHECK-NOT: call void @llvm.memcpy
 
 // CHECK-NOT: call void @llvm.memcpy

tests/codegen/mem-replace-direct-memcpy.rs

@@ -0,0 +1,34 @@
+// compile-flags: -O -C no-prepopulate-passes
+// min-llvm-version: 15.0 (for opaque pointers)
+// only-x86_64 (to not worry about usize differing)
+// ignore-debug (the debug assertions get in the way)
+
+#![crate_type = "lib"]
+
+#[no_mangle]
+// CHECK-LABEL: @replace_usize(
+pub fn replace_usize(r: &mut usize, v: usize) -> usize {
+    // CHECK-NOT: alloca
+    // CHECK: %[[R:.+]] = load i64, ptr %r
+    // CHECK: store i64 %v, ptr %r
+    // CHECK: ret i64 %[[R]]
+    std::mem::replace(r, v)
+}
+
+#[no_mangle]
+// CHECK-LABEL: @replace_ref_str(
+pub fn replace_ref_str<'a>(r: &mut &'a str, v: &'a str) -> &'a str {
+    // CHECK-NOT: alloca
+    // CHECK: %[[A:.+]] = load ptr
+    // CHECK: %[[B:.+]] = load i64
+    // CHECK-NOT: store
+    // CHECK-NOT: load
+    // CHECK: store ptr
+    // CHECK: store i64
+    // CHECK-NOT: load
+    // CHECK-NOT: store
+    // CHECK: %[[P1:.+]] = insertvalue { ptr, i64 } poison, ptr %[[A]], 0
+    // CHECK: %[[P2:.+]] = insertvalue { ptr, i64 } %[[P1]], i64 %[[B]], 1
+    // CHECK: ret { ptr, i64 } %[[P2]]
+    std::mem::replace(r, v)
+}