replace box_new in Box::new with write_via_move

requires lowering write_via_move during MIR building to make it just like an assignment

Author: RalfJung

compiler/rustc_mir_build/src/builder/expr/into.rs

@@ -9,8 +9,8 @@ use rustc_middle::mir::*;
 use rustc_middle::span_bug;
 use rustc_middle::thir::*;
 use rustc_middle::ty::{self, CanonicalUserTypeAnnotation, Ty};
-use rustc_span::DUMMY_SP;
 use rustc_span::source_map::Spanned;
+use rustc_span::{DUMMY_SP, sym};
 use rustc_trait_selection::infer::InferCtxtExt;
 use tracing::{debug, instrument};
 
@@ -365,6 +365,31 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                     None
                 })
             }
+            // The `write_via_move` intrinsic needs to be special-cased very early to avoid
+            // introducing unnecessary copies that can be hard to remove again later:
+            // `write_via_move(ptr, val)` becomes `*ptr = val` but without any dropping.
+            ExprKind::Call { ty, fun, ref args, .. }
+                if let ty::FnDef(def_id, _generic_args) = ty.kind()
+                    && let Some(intrinsic) = this.tcx.intrinsic(def_id)
+                    && intrinsic.name == sym::write_via_move =>
+            {
+                // We still have to evaluate the callee expression as normal (but we don't care
+                // about its result).
+                let _fun = unpack!(block = this.as_local_operand(block, fun));
+                // The destination must have unit type (so we don't actually have to store anything
+                // into it).
+                assert!(destination.ty(&this.local_decls, this.tcx).ty.is_unit());
+
+                // Compile this to an assignment of the argument into the destination.
+                let [ptr, val] = **args else {
+                    span_bug!(expr_span, "invalid write_via_move call")
+                };
+                let Some(ptr) = unpack!(block = this.as_local_operand(block, ptr)).place() else {
+                    span_bug!(expr_span, "invalid write_via_move call")
+                };
+                let ptr_deref = ptr.project_deeper(&[ProjectionElem::Deref], this.tcx);
+                this.expr_into_dest(ptr_deref, block, val)
+            }
             ExprKind::Call { ty: _, fun, ref args, from_hir_call, fn_span } => {
                 let fun = unpack!(block = this.as_local_operand(block, fun));
                 let args: Box<[_]> = args

compiler/rustc_mir_transform/src/lower_intrinsics.rs

@@ -171,30 +171,7 @@ impl<'tcx> crate::MirPass<'tcx> for LowerIntrinsics {
                             Some(target) => TerminatorKind::Goto { target },
                         }
                     }
-                    sym::write_via_move => {
-                        let target = target.unwrap();
-                        let Ok([ptr, val]) = take_array(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.node.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::new(
-                            terminator.source_info,
-                            StatementKind::Assign(Box::new((derefed_place, Rvalue::Use(val.node)))),
-                        ));
-                        terminator.kind = TerminatorKind::Goto { target };
-                    }
+                    // `write_via_move` is already lowered during MIR building.
                     sym::discriminant_value => {
                         let target = target.unwrap();
                         let Ok([arg]) = take_array(args) else {

library/alloc/src/alloc.rs

@@ -343,11 +343,15 @@ unsafe impl Allocator for Global {
 }
 
 /// The allocator for `Box`.
+///
+/// # Safety
+///
+/// `size` and `align` must satisfy the conditions in [`Layout::from_size_align`].
 #[cfg(not(no_global_oom_handling))]
 #[lang = "exchange_malloc"]
 #[inline]
 #[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
-unsafe fn exchange_malloc(size: usize, align: usize) -> *mut u8 {
+pub(crate) unsafe fn exchange_malloc(size: usize, align: usize) -> *mut u8 {
     let layout = unsafe { Layout::from_size_align_unchecked(size, align) };
     match Global.allocate(layout) {
         Ok(ptr) => ptr.as_mut_ptr(),

library/alloc/src/boxed.rs

@@ -203,7 +203,7 @@ use core::task::{Context, Poll};
 
 #[cfg(not(no_global_oom_handling))]
 use crate::alloc::handle_alloc_error;
-use crate::alloc::{AllocError, Allocator, Global, Layout};
+use crate::alloc::{AllocError, Allocator, Global, Layout, exchange_malloc};
 use crate::raw_vec::RawVec;
 #[cfg(not(no_global_oom_handling))]
 use crate::str::from_boxed_utf8_unchecked;
@@ -259,7 +259,15 @@ impl<T> Box<T> {
     #[rustc_diagnostic_item = "box_new"]
     #[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
     pub fn new(x: T) -> Self {
-        return box_new(x);
+        // SAFETY: the size and align of a valid type `T` are always valid for `Layout`.
+        let ptr = unsafe {
+            exchange_malloc(<T as SizedTypeProperties>::SIZE, <T as SizedTypeProperties>::ALIGN)
+        } as *mut T;
+        // Nothing below can panic so we do not have to worry about deallocating `ptr`.
+        // SAFETY: we just allocated the box to store `x`.
+        unsafe { core::intrinsics::write_via_move(ptr, x) };
+        // SAFETY: we just initialized `b`.
+        unsafe { mem::transmute(ptr) }
     }
 
     /// Constructs a new box with uninitialized contents.

tests/mir-opt/building/write_via_move.box_new.CleanupPostBorrowck.after.mir

@@ -0,0 +1,76 @@
+// MIR for `box_new` after CleanupPostBorrowck
+
+fn box_new(_1: T) -> Box<[T; 1024]> {
+    debug x => _1;
+    let mut _0: std::boxed::Box<[T; 1024]>;
+    let mut _2: std::boxed::Box<std::mem::MaybeUninit<[T; 1024]>>;
+    let mut _4: &mut std::mem::MaybeUninit<[T; 1024]>;
+    let mut _5: &mut std::mem::MaybeUninit<[T; 1024]>;
+    let _6: ();
+    let mut _7: *mut [T; 1024];
+    let mut _8: T;
+    let mut _9: std::boxed::Box<std::mem::MaybeUninit<[T; 1024]>>;
+    scope 1 {
+        debug b => _2;
+        let _3: *mut [T; 1024];
+        scope 2 {
+            debug ptr => _3;
+        }
+    }
+
+    bb0: {
+        StorageLive(_2);
+        _2 = Box::<[T; 1024]>::new_uninit() -> [return: bb1, unwind: bb7];
+    }
+
+    bb1: {
+        nop;
+        StorageLive(_3);
+        StorageLive(_4);
+        StorageLive(_5);
+        _5 = &mut (*_2);
+        _4 = &mut (*_5);
+        _3 = MaybeUninit::<[T; 1024]>::as_mut_ptr(move _4) -> [return: bb2, unwind: bb6];
+    }
+
+    bb2: {
+        StorageDead(_4);
+        nop;
+        StorageDead(_5);
+        StorageLive(_6);
+        StorageLive(_7);
+        _7 = copy _3;
+        StorageLive(_8);
+        _8 = copy _1;
+        (*_7) = [move _8; 1024];
+        StorageDead(_8);
+        StorageDead(_7);
+        StorageDead(_6);
+        StorageLive(_9);
+        _9 = move _2;
+        _0 = Box::<MaybeUninit<[T; 1024]>>::assume_init(move _9) -> [return: bb3, unwind: bb5];
+    }
+
+    bb3: {
+        StorageDead(_9);
+        StorageDead(_3);
+        drop(_2) -> [return: bb4, unwind: bb7];
+    }
+
+    bb4: {
+        StorageDead(_2);
+        return;
+    }
+
+    bb5 (cleanup): {
+        drop(_9) -> [return: bb6, unwind terminate(cleanup)];
+    }
+
+    bb6 (cleanup): {
+        drop(_2) -> [return: bb7, unwind terminate(cleanup)];
+    }
+
+    bb7 (cleanup): {
+        resume;
+    }
+}

tests/mir-opt/building/write_via_move.rs

@@ -0,0 +1,23 @@
+//! Ensure we don't generate unnecessary copys for `write_via_move`.
+//@ compile-flags: -Zmir-opt-level=0
+#![feature(core_intrinsics)]
+
+use std::mem;
+
+// Can't emit `built.after` here as that contains user type annotations which contain DefId that
+// change all the time.
+// EMIT_MIR write_via_move.box_new.CleanupPostBorrowck.after.mir
+// CHECK-LABEL: fn box_new
+#[inline(never)]
+fn box_new<T: Copy>(x: T) -> Box<[T; 1024]> {
+    let mut b = Box::new_uninit();
+    let ptr = mem::MaybeUninit::as_mut_ptr(&mut *b);
+    // Ensure the array gets constructed directly into the deref'd pointer.
+    // CHECK: (*[[TEMP1:_.+]]) = [{{(move|copy) _.+}}; 1024];
+    unsafe { std::intrinsics::write_via_move(ptr, [x; 1024]) };
+    unsafe { b.assume_init() }
+}
+
+fn main() {
+    box_new(0);
+}

tests/mir-opt/inline_coroutine_body.run2-{closure#0}.Inline.panic-unwind.diff

@@ -213,7 +213,7 @@
 +         StorageLive(_42);
 +         _42 = Option::<()>::None;
 +         _35 = copy ((*_37).0: std::option::Option<()>);
-+         ((*_37).0: std::option::Option<()>) = copy _42;
++         ((*_37).0: std::option::Option<()>) = move _42;
 +         StorageDead(_42);
 +         StorageLive(_43);
 +         _43 = discriminant(_35);

tests/mir-opt/lower_intrinsics.rs

@@ -197,17 +197,6 @@ 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) {
-    // CHECK-LABEL: fn write_via_move_string(
-    // CHECK: [[ptr:_.*]] = &raw mut (*_1);
-    // CHECK: [[tmp:_.*]] = move _2;
-    // CHECK: (*[[ptr]]) = move [[tmp]];
-    // CHECK: return;
-
-    unsafe { core::intrinsics::write_via_move(r, v) }
-}
-
 pub enum Never {}
 
 // EMIT_MIR lower_intrinsics.ptr_offset.LowerIntrinsics.diff