From 6841b6119c5bd0a687c004d872b45333c9a017ac Mon Sep 17 00:00:00 2001
From: Benno Lossin <benno.lossin@proton.me>
Date: Mon, 19 Aug 2024 11:24:39 +0000
Subject: [PATCH] rust: init: add `write_[pin_]init` functions

Sometimes it is necessary to split allocation and initialization into
two steps. One such situation is when reusing existing allocations
obtained via `Box::drop_contents`. See [1] for an example.

In order to support this use case add `write_[pin_]init` functions to the
pin-init API. These functions operate on already allocated smart
pointers that wrap `MaybeUninit<T>`.

Link: https://lore.kernel.org/rust-for-linux/f026532f-8594-4f18-9aa5-57ad3f5bc592@proton.me/ [1]
Signed-off-by: Benno Lossin <benno.lossin@proton.me>
Reviewed-by: Boqun Feng <boqun.feng@gmail.com>
Reviewed-by: Alice Ryhl <aliceryhl@google.com>
Reviewed-by: Gary Guo <gary@garyguo.net>
Link: https://lore.kernel.org/r/20240819112415.99810-2-benno.lossin@proton.me
Signed-off-by: Miguel Ojeda <ojeda@kernel.org>
Signed-off-by: Paolo Bonzini <pbonzini@redhat.com>
(cherry picked from commit 6d1c22d0ace31d096b0dab5318c6a0d3219d6456)
---
 src/lib.rs | 88 ++++++++++++++++++++++++++++++++++++++----------------
 1 file changed, 63 insertions(+), 25 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 186a2fb..81701c7 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1137,13 +1137,7 @@ impl<T> InPlaceInit<T> for Box<T> {
     where
         E: From<AllocError>,
     {
-        let mut this = Box::try_new_uninit()?;
-        let slot = this.as_mut_ptr();
-        // SAFETY: When init errors/panics, slot will get deallocated but not dropped,
-        // slot is valid and will not be moved, because we pin it later.
-        unsafe { init.__pinned_init(slot)? };
-        // SAFETY: All fields have been initialized.
-        Ok(unsafe { this.assume_init() }.into())
+        Box::try_new_uninit()?.write_pin_init(init)
     }
 
     #[inline]
@@ -1151,13 +1145,7 @@ impl<T> InPlaceInit<T> for Box<T> {
     where
         E: From<AllocError>,
     {
-        let mut this = Box::try_new_uninit()?;
-        let slot = this.as_mut_ptr();
-        // SAFETY: When init errors/panics, slot will get deallocated but not dropped,
-        // slot is valid.
-        unsafe { init.__init(slot)? };
-        // SAFETY: All fields have been initialized.
-        Ok(unsafe { this.assume_init() })
+        Box::try_new_uninit()?.write_init(init)
     }
 }
 
@@ -1168,14 +1156,7 @@ impl<T> InPlaceInit<T> for Arc<T> {
     where
         E: From<AllocError>,
     {
-        let mut this = Arc::try_new_uninit()?;
-        let slot = unsafe { Arc::get_mut_unchecked(&mut this) };
-        let slot = slot.as_mut_ptr();
-        // SAFETY: When init errors/panics, slot will get deallocated but not dropped,
-        // slot is valid and will not be moved, because we pin it later.
-        unsafe { init.__pinned_init(slot)? };
-        // SAFETY: All fields have been initialized and this is the only `Arc` to that data.
-        Ok(unsafe { Pin::new_unchecked(this.assume_init()) })
+        Arc::try_new_uninit()?.write_pin_init(init)
     }
 
     #[inline]
@@ -1183,14 +1164,71 @@ impl<T> InPlaceInit<T> for Arc<T> {
     where
         E: From<AllocError>,
     {
-        let mut this = Arc::try_new_uninit()?;
-        let slot = unsafe { Arc::get_mut_unchecked(&mut this) };
+        Arc::try_new_uninit()?.write_init(init)
+    }
+}
+
+/// Smart pointer containing uninitialized memory and that can write a value.
+pub trait InPlaceWrite<T> {
+    /// The type `Self` turns into when the contents are initialized.
+    type Initialized;
+
+    /// Use the given initializer to write a value into `self`.
+    ///
+    /// Does not drop the current value and considers it as uninitialized memory.
+    fn write_init<E>(self, init: impl Init<T, E>) -> Result<Self::Initialized, E>;
+
+    /// Use the given pin-initializer to write a value into `self`.
+    ///
+    /// Does not drop the current value and considers it as uninitialized memory.
+    fn write_pin_init<E>(self, init: impl PinInit<T, E>) -> Result<Pin<Self::Initialized>, E>;
+}
+
+#[cfg(feature = "alloc")]
+impl<T> InPlaceWrite<T> for Box<MaybeUninit<T>> {
+    type Initialized = Box<T>;
+
+    fn write_init<E>(mut self, init: impl Init<T, E>) -> Result<Self::Initialized, E> {
+        let slot = self.as_mut_ptr();
+        // SAFETY: When init errors/panics, slot will get deallocated but not dropped,
+        // slot is valid.
+        unsafe { init.__init(slot)? };
+        // SAFETY: All fields have been initialized.
+        Ok(unsafe { self.assume_init() })
+    }
+
+    fn write_pin_init<E>(mut self, init: impl PinInit<T, E>) -> Result<Pin<Self::Initialized>, E> {
+        let slot = self.as_mut_ptr();
+        // SAFETY: When init errors/panics, slot will get deallocated but not dropped,
+        // slot is valid and will not be moved, because we pin it later.
+        unsafe { init.__pinned_init(slot)? };
+        // SAFETY: All fields have been initialized.
+        Ok(unsafe { self.assume_init() }.into())
+    }
+}
+
+#[cfg(feature = "alloc")]
+impl<T> InPlaceWrite<T> for Arc<MaybeUninit<T>> {
+    type Initialized = Arc<T>;
+
+    fn write_init<E>(mut self, init: impl Init<T, E>) -> Result<Self::Initialized, E> {
+        let slot = unsafe { Arc::get_mut_unchecked(&mut self) };
         let slot = slot.as_mut_ptr();
         // SAFETY: When init errors/panics, slot will get deallocated but not dropped,
         // slot is valid.
         unsafe { init.__init(slot)? };
         // SAFETY: All fields have been initialized.
-        Ok(unsafe { this.assume_init() })
+        Ok(unsafe { self.assume_init() })
+    }
+
+    fn write_pin_init<E>(mut self, init: impl PinInit<T, E>) -> Result<Pin<Self::Initialized>, E> {
+        let slot = unsafe { Arc::get_mut_unchecked(&mut self) };
+        let slot = slot.as_mut_ptr();
+        // SAFETY: When init errors/panics, slot will get deallocated but not dropped,
+        // slot is valid and will not be moved, because we pin it later.
+        unsafe { init.__pinned_init(slot)? };
+        // SAFETY: All fields have been initialized and this is the only `Arc` to that data.
+        Ok(unsafe { Pin::new_unchecked(self.assume_init()) })
     }
 }