From 2e02a9ead61ccc1837572cc961d38b9be0592d8f Mon Sep 17 00:00:00 2001
From: Martin Habovstiak <martin.habovstiak@gmail.com>
Date: Fri, 4 Feb 2022 17:53:32 +0100
Subject: [PATCH] Added `Box::take()` method

The new `Box::take()` method allows to read from `Box` and reuse the
allocation safely. This commit also changes some `unsafe` code in the
compiler to use this method instead and becomes safe.
---
 compiler/rustc_data_structures/src/functor.rs | 12 +---
 library/alloc/src/boxed.rs                    | 56 +++++++++++++++++++
 library/alloc/src/lib.rs                      |  1 +
 3 files changed, 59 insertions(+), 10 deletions(-)

diff --git a/compiler/rustc_data_structures/src/functor.rs b/compiler/rustc_data_structures/src/functor.rs
index a3d3f988344c6..0fe3318e89447 100644
--- a/compiler/rustc_data_structures/src/functor.rs
+++ b/compiler/rustc_data_structures/src/functor.rs
@@ -17,16 +17,8 @@ impl<T> IdFunctor for Box<T> {
     where
         F: FnMut(Self::Inner) -> Result<Self::Inner, E>,
     {
-        let raw = Box::into_raw(self);
-        Ok(unsafe {
-            // SAFETY: The raw pointer points to a valid value of type `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 raw: Box<mem::MaybeUninit<T>> = Box::from_raw(raw.cast());
-            // SAFETY: Write the mapped value back into the `Box`.
-            Box::write(raw, f(value)?)
-        })
+        let (value, allocation) = Box::take(self);
+        Ok(Box::write(allocation, f(value)?))
     }
 }
 
diff --git a/library/alloc/src/boxed.rs b/library/alloc/src/boxed.rs
index aa7344ba405a9..bfee9492f4369 100644
--- a/library/alloc/src/boxed.rs
+++ b/library/alloc/src/boxed.rs
@@ -578,6 +578,62 @@ impl<T, A: Allocator> Box<T, A> {
     {
         *boxed
     }
+
+    /// Safely reads the value on heap without deallocating.
+    ///
+    /// This method is mostly useful for avoiding reallocations.
+    /// The method consumes the box to prevent double-free.
+    /// The value on heap should be considered truly uninitialized and MUST NOT be read.
+    /// Specifically, this code is definitely UB:
+    ///
+    /// ```no_run
+    /// #![feature(new_uninit)]
+    ///
+    /// let value = Box::new("hello".to_owned());
+    /// // Double free here!
+    /// unsafe { Box::take(value).1.assume_init(); }
+    /// ```
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(new_uninit)]
+    ///
+    /// let boxed_value = Box::new("hello".to_owned());
+    /// let (mut value, allocation) = Box::take(boxed_value);
+    /// assert_eq!(value, "hello");
+    /// // More realistic code would consume the value and produce other value.
+    /// // For simple demonstration we just modify it here.
+    /// value.push_str(" world");
+    /// let boxed_value = Box::write(allocation, value);
+    /// assert_eq!(*boxed_value, "hello world");
+    /// ```
+    #[unstable(feature = "new_uninit", issue = "63291")]
+    #[rustc_const_unstable(feature = "const_box", issue = "92521")]
+    #[inline]
+    pub const fn take(this: Self) -> (T, Box<mem::MaybeUninit<T>, A>) {
+        let (raw, alloc) = Box::into_raw_with_allocator(this);
+        unsafe {
+            // SAFETY:
+            // * The pointer given to from_raw_in was obtained from box above using the same
+            //   allocator
+            // * Casting `*mut T` to `*mut MaybeUninit<T>` is sound because they have same layout
+            // * Safely obtaining `&mut T` pointing into the allocation will become impossible after
+            //   this call because we consume `this` (no variance issues)
+            // * `MaybeUninit` disables destructor of `T` so it can't double free
+            // * Leak shouldn't happen because `assume_init_read` below doesn't panic and the value
+            //   can be dropped afterwards
+            let new_box = Box::from_raw_in(raw.cast::<mem::MaybeUninit<T>>(), alloc);
+            // SAFETY:
+            // * The value being read was proven to be initialized when this function was called
+            // * We didn't touch the value inside this function, just cast the pointer, so it's
+            //   still valid.
+            // * We don't read the value on heap again and prevent safe code from reading it by
+            //   marking the box with `MaybeUninit`.
+            let value = new_box.assume_init_read();
+            (value, new_box)
+        }
+    }
 }
 
 impl<T> Box<[T]> {
diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs
index dfd3771c1d042..8aaa1312694c4 100644
--- a/library/alloc/src/lib.rs
+++ b/library/alloc/src/lib.rs
@@ -101,6 +101,7 @@
 #![feature(const_size_of_val)]
 #![feature(const_align_of_val)]
 #![feature(const_ptr_read)]
+#![feature(const_maybe_uninit_assume_init_read)]
 #![feature(const_maybe_uninit_write)]
 #![feature(const_maybe_uninit_as_mut_ptr)]
 #![feature(const_refs_to_cell)]