From 41e21aa1c27522c6b29c5b88e37b1f479f63e38c Mon Sep 17 00:00:00 2001
From: Martin Habovstiak <martin.habovstiak@gmail.com>
Date: Mon, 13 Sep 2021 15:44:27 +0200
Subject: [PATCH] Implement write() method for Box<MaybeUninit<T>>

This adds method similar to `MaybeUninit::write` main difference being
it returns owned `Box`. This can be used to elide copy from stack
safely, however it's not currently tested that the optimization actually
occurs.

Analogous methods are not provided for `Rc` and `Arc` as those need to
handle the possibility of sharing. Some version of them may be added in
the future.

This was discussed in #63291 which this change extends.
---
 compiler/rustc_data_structures/src/functor.rs |  6 ++--
 library/alloc/src/boxed.rs                    | 36 +++++++++++++++++++
 2 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/compiler/rustc_data_structures/src/functor.rs b/compiler/rustc_data_structures/src/functor.rs
index 920f7b1ed0a54..2703af97a1be5 100644
--- a/compiler/rustc_data_structures/src/functor.rs
+++ b/compiler/rustc_data_structures/src/functor.rs
@@ -31,11 +31,9 @@ impl<T> IdFunctor for Box<T> {
             let value = raw.read();
             // SAFETY: Converts `Box<T>` to `Box<MaybeUninit<T>>` which is the
             // inverse of `Box::assume_init()` and should be safe.
-            let mut raw: Box<mem::MaybeUninit<T>> = Box::from_raw(raw.cast());
+            let raw: Box<mem::MaybeUninit<T>> = Box::from_raw(raw.cast());
             // SAFETY: Write the mapped value back into the `Box`.
-            raw.write(f(value)?);
-            // SAFETY: We just initialized `raw`.
-            raw.assume_init()
+            Box::write(raw, f(value)?)
         })
     }
 }
diff --git a/library/alloc/src/boxed.rs b/library/alloc/src/boxed.rs
index f6332b072cf30..0b72b3f0ee7b0 100644
--- a/library/alloc/src/boxed.rs
+++ b/library/alloc/src/boxed.rs
@@ -763,6 +763,42 @@ impl<T, A: Allocator> Box<mem::MaybeUninit<T>, A> {
         let (raw, alloc) = Box::into_raw_with_allocator(self);
         unsafe { Box::from_raw_in(raw as *mut T, alloc) }
     }
+
+    /// Writes the value and converts to `Box<T, A>`.
+    ///
+    /// This method converts the box similarly to [`Box::assume_init`] but
+    /// writes `value` into it before conversion thus guaranteeing safety.
+    /// In some scenarios use of this method may improve performance because
+    /// the compiler may be able to optimize copying from stack.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(new_uninit)]
+    ///
+    /// let big_box = Box::<[usize; 1024]>::new_uninit();
+    ///
+    /// let mut array = [0; 1024];
+    /// for (i, place) in array.iter_mut().enumerate() {
+    ///     *place = i;
+    /// }
+    ///
+    /// // The optimizer may be able to elide this copy, so previous code writes
+    /// // to heap directly.
+    /// let big_box = Box::write(big_box, array);
+    ///
+    /// for (i, x) in big_box.iter().enumerate() {
+    ///     assert_eq!(*x, i);
+    /// }
+    /// ```
+    #[unstable(feature = "new_uninit", issue = "63291")]
+    #[inline]
+    pub fn write(mut boxed: Self, value: T) -> Box<T, A> {
+        unsafe {
+            (*boxed).write(value);
+            boxed.assume_init()
+        }
+    }
 }
 
 impl<T, A: Allocator> Box<[mem::MaybeUninit<T>], A> {