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

bors · bors · commit 6db1e5e771cb · 2023-05-01T14:29:15.000Z
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.)
diff --git a/compiler/rustc_hir_analysis/src/check/intrinsic.rs b/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(
diff --git a/compiler/rustc_mir_transform/src/lower_intrinsics.rs b/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);
diff --git a/compiler/rustc_span/src/symbol.rs b/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,
diff --git a/library/core/src/intrinsics.rs b/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);
+    }
+}
diff --git a/library/core/src/ptr/mod.rs b/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)
     }
 }
 
diff --git a/src/tools/miri/tests/fail/dangling_pointers/null_pointer_write_zst.rs b/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
 }
diff --git a/src/tools/miri/tests/fail/dangling_pointers/null_pointer_write_zst.stderr b/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
diff --git a/tests/codegen/mem-replace-big-type.rs b/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
diff --git a/tests/codegen/mem-replace-direct-memcpy.rs b/tests/codegen/mem-replace-direct-memcpy.rs
diff --git a/tests/codegen/mem-replace-simple-type.rs b/tests/codegen/mem-replace-simple-type.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)
+}
diff --git a/tests/mir-opt/lower_intrinsics.option_payload.LowerIntrinsics.diff b/tests/mir-opt/lower_intrinsics.option_payload.LowerIntrinsics.diff
@@ -24,7 +24,7 @@
           _4 = &raw const (*_1);           // scope 1 at $DIR/lower_intrinsics.rs:+2:55: +2:56
 -         _3 = option_payload_ptr::<usize>(move _4) -> [return: bb1, unwind unreachable]; // scope 1 at $DIR/lower_intrinsics.rs:+2:18: +2:57
 -                                          // mir::Constant
--                                          // + span: $DIR/lower_intrinsics.rs:132:18: 132:54
+-                                          // + span: $DIR/lower_intrinsics.rs:137:18: 137:54
 -                                          // + literal: Const { ty: unsafe extern "rust-intrinsic" fn(*const Option<usize>) -> *const usize {option_payload_ptr::<usize>}, val: Value(<ZST>) }
 +         _3 = &raw const (((*_4) as Some).0: usize); // scope 1 at $DIR/lower_intrinsics.rs:+2:18: +2:57
 +         goto -> bb1;                     // scope 1 at $DIR/lower_intrinsics.rs:+2:18: +2:57
@@ -37,7 +37,7 @@
           _6 = &raw const (*_2);           // scope 2 at $DIR/lower_intrinsics.rs:+3:55: +3:56
 -         _5 = option_payload_ptr::<String>(move _6) -> [return: bb2, unwind unreachable]; // scope 2 at $DIR/lower_intrinsics.rs:+3:18: +3:57
 -                                          // mir::Constant
--                                          // + span: $DIR/lower_intrinsics.rs:133:18: 133:54
+-                                          // + span: $DIR/lower_intrinsics.rs:138:18: 138:54
 -                                          // + literal: Const { ty: unsafe extern "rust-intrinsic" fn(*const Option<String>) -> *const String {option_payload_ptr::<String>}, val: Value(<ZST>) }
 +         _5 = &raw const (((*_6) as Some).0: std::string::String); // scope 2 at $DIR/lower_intrinsics.rs:+3:18: +3:57
 +         goto -> bb2;                     // scope 2 at $DIR/lower_intrinsics.rs:+3:18: +3:57
diff --git a/tests/mir-opt/lower_intrinsics.ptr_offset.LowerIntrinsics.diff b/tests/mir-opt/lower_intrinsics.ptr_offset.LowerIntrinsics.diff
@@ -15,7 +15,7 @@
           _4 = _2;                         // scope 0 at $DIR/lower_intrinsics.rs:+1:33: +1:34
 -         _0 = offset::<*const i32, isize>(move _3, move _4) -> [return: bb1, unwind unreachable]; // scope 0 at $DIR/lower_intrinsics.rs:+1:5: +1:35
 -                                          // mir::Constant
--                                          // + span: $DIR/lower_intrinsics.rs:139:5: 139:29
+-                                          // + span: $DIR/lower_intrinsics.rs:144:5: 144:29
 -                                          // + literal: Const { ty: unsafe extern "rust-intrinsic" fn(*const i32, isize) -> *const i32 {offset::<*const i32, isize>}, val: Value(<ZST>) }
 +         _0 = Offset(move _3, move _4);   // scope 0 at $DIR/lower_intrinsics.rs:+1:5: +1:35
 +         goto -> bb1;                     // scope 0 at $DIR/lower_intrinsics.rs:+1:5: +1:35
diff --git a/tests/mir-opt/lower_intrinsics.rs b/tests/mir-opt/lower_intrinsics.rs
@@ -124,6 +124,11 @@ pub fn read_via_copy_uninhabited(r: &Never) -> Never {
     unsafe { core::intrinsics::read_via_copy(r) }
 }
 
+// EMIT_MIR lower_intrinsics.write_via_move_string.LowerIntrinsics.diff
+pub fn write_via_move_string(r: &mut String, v: String) {
+    unsafe { core::intrinsics::write_via_move(r, v) }
+}
+
 pub enum Never {}
 
 // EMIT_MIR lower_intrinsics.option_payload.LowerIntrinsics.diff
diff --git a/tests/mir-opt/lower_intrinsics.write_via_move_string.LowerIntrinsics.diff b/tests/mir-opt/lower_intrinsics.write_via_move_string.LowerIntrinsics.diff
@@ -0,0 +1,36 @@
+- // MIR for `write_via_move_string` before LowerIntrinsics
++ // MIR for `write_via_move_string` after LowerIntrinsics
+  
+  fn write_via_move_string(_1: &mut String, _2: String) -> () {
+      debug r => _1;                       // in scope 0 at $DIR/lower_intrinsics.rs:+0:30: +0:31
+      debug v => _2;                       // in scope 0 at $DIR/lower_intrinsics.rs:+0:46: +0:47
+      let mut _0: ();                      // return place in scope 0 at $DIR/lower_intrinsics.rs:+0:57: +0:57
+      let mut _3: *mut std::string::String; // in scope 0 at $DIR/lower_intrinsics.rs:+1:47: +1:48
+      let mut _4: std::string::String;     // in scope 0 at $DIR/lower_intrinsics.rs:+1:50: +1:51
+      scope 1 {
+      }
+  
+      bb0: {
+          StorageLive(_3);                 // scope 1 at $DIR/lower_intrinsics.rs:+1:47: +1:48
+          _3 = &raw mut (*_1);             // scope 1 at $DIR/lower_intrinsics.rs:+1:47: +1:48
+          StorageLive(_4);                 // scope 1 at $DIR/lower_intrinsics.rs:+1:50: +1:51
+          _4 = move _2;                    // scope 1 at $DIR/lower_intrinsics.rs:+1:50: +1:51
+-         _0 = write_via_move::<String>(move _3, move _4) -> [return: bb1, unwind unreachable]; // scope 1 at $DIR/lower_intrinsics.rs:+1:14: +1:52
+-                                          // mir::Constant
+-                                          // + span: $DIR/lower_intrinsics.rs:129:14: 129:46
+-                                          // + literal: Const { ty: unsafe extern "rust-intrinsic" fn(*mut String, String) {write_via_move::<String>}, val: Value(<ZST>) }
++         (*_3) = move _4;                 // scope 1 at $DIR/lower_intrinsics.rs:+1:14: +1:52
++         goto -> bb1;                     // scope 1 at $DIR/lower_intrinsics.rs:+1:14: +1:52
+      }
+  
+      bb1: {
+          StorageDead(_4);                 // scope 1 at $DIR/lower_intrinsics.rs:+1:51: +1:52
+          StorageDead(_3);                 // scope 1 at $DIR/lower_intrinsics.rs:+1:51: +1:52
+          goto -> bb2;                     // scope 0 at $DIR/lower_intrinsics.rs:+2:1: +2:2
+      }
+  
+      bb2: {
+          return;                          // scope 0 at $DIR/lower_intrinsics.rs:+2:2: +2:2
+      }
+  }
+  
diff --git a/tests/mir-opt/pre-codegen/mem_replace.manual_replace.PreCodegen.after.mir b/tests/mir-opt/pre-codegen/mem_replace.manual_replace.PreCodegen.after.mir
@@ -0,0 +1,16 @@
+// MIR for `manual_replace` after PreCodegen
+
+fn manual_replace(_1: &mut u32, _2: u32) -> u32 {
+    debug r => _1;                       // in scope 0 at $DIR/mem_replace.rs:+0:23: +0:24
+    debug v => _2;                       // in scope 0 at $DIR/mem_replace.rs:+0:36: +0:37
+    let mut _0: u32;                     // return place in scope 0 at $DIR/mem_replace.rs:+1:9: +1:13
+    scope 1 {
+        debug temp => _0;                // in scope 1 at $DIR/mem_replace.rs:+1:9: +1:13
+    }
+
+    bb0: {
+        _0 = (*_1);                      // scope 0 at $DIR/mem_replace.rs:+1:16: +1:18
+        (*_1) = _2;                      // scope 1 at $DIR/mem_replace.rs:+2:5: +2:11
+        return;                          // scope 0 at $DIR/mem_replace.rs:+4:2: +4:2
+    }
+}
diff --git a/tests/mir-opt/pre-codegen/mem_replace.mem_replace.PreCodegen.after.mir b/tests/mir-opt/pre-codegen/mem_replace.mem_replace.PreCodegen.after.mir
diff --git a/tests/mir-opt/pre-codegen/mem_replace.rs b/tests/mir-opt/pre-codegen/mem_replace.rs

Original file line number	Diff line number	Diff line change
`@@ -7,5 +7,5 @@ fn main() {`
`7`	`7`	`// Also not assigning directly as that's array initialization, not assignment.`
`8`	`8`	`let zst_val = [1u8; 0];`
`9`	`9`	`unsafe { std::ptr::null_mut::<[u8; 0]>().write(zst_val) };`
`10`		`- //~^ERROR: memory access failed: null pointer is a dangling pointer`
	`10`	`+ //~^ERROR: dereferencing pointer failed: null pointer is a dangling pointer`
`11`	`11`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1,8 +1,8 @@`
`1`		`-error: Undefined Behavior: memory access failed: null pointer is a dangling pointer (it has no provenance)`
	`1`	`+error: Undefined Behavior: dereferencing pointer failed: null pointer is a dangling pointer (it has no provenance)`
`2`	`2`	`--> $DIR/null_pointer_write_zst.rs:LL:CC`
`3`	`3`	`\|`
`4`	`4`	`LL \| unsafe { std::ptr::null_mut::<[u8; 0]>().write(zst_val) };`
`5`		`- \| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ memory access failed: null pointer is a dangling pointer (it has no provenance)`
	`5`	`+ \| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ dereferencing pointer failed: null pointer is a dangling pointer (it has no provenance)`
`6`	`6`	`\|`
`7`	`7`	`= help: this indicates a bug in the program: it performed an invalid operation, and caused Undefined Behavior`
`8`	`8`	`= help: see https://doc.rust-lang.org/nightly/reference/behavior-considered-undefined.html for further information`