Auto merge of #79607 - DrMeepster:maybe_uninit_write_slice, r=m-ou-se

MaybeUninit::copy/clone_from_slice

This PR adds 2 new methods to MaybeUninit under the feature of `maybe_uninit_write_slice`: `copy_from_slice` and `clone_from_slice`.

These are useful for initializing uninitialized buffers (such as the one returned by `Vec::spare_capacity_mut` for example) with initialized data.

The methods behave similarly to the methods on slices, but the destination is uninitialized and they return the destination slice as an initialized slice.

3 files changed, 276 insertions, 0 deletions

diff --git a/library/core/src/mem/maybe_uninit.rs b/library/core/src/mem/maybe_uninit.rs
index 8800d7714cf..57e0bb1499b 100644
--- a/library/core/src/mem/maybe_uninit.rs
+++ b/library/core/src/mem/maybe_uninit.rs
@@ -860,4 +860,154 @@ impl<T> MaybeUninit<T> {
     pub const fn slice_as_mut_ptr(this: &mut [MaybeUninit<T>]) -> *mut T {
         this.as_mut_ptr() as *mut T
     }
+
+    /// Copies the elements from `src` to `this`, returning a mutable reference to the now initalized contents of `this`.
+    ///
+    /// If `T` does not implement `Copy`, use [`write_slice_cloned`]
+    ///
+    /// This is similar to [`slice::copy_from_slice`].
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if the two slices have different lengths.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(maybe_uninit_write_slice)]
+    /// use std::mem::MaybeUninit;
+    ///
+    /// let mut dst = [MaybeUninit::uninit(); 32];
+    /// let src = [0; 32];
+    ///
+    /// let init = MaybeUninit::write_slice(&mut dst, &src);
+    ///
+    /// assert_eq!(init, src);
+    /// ```
+    ///
+    /// ```
+    /// #![feature(maybe_uninit_write_slice, vec_spare_capacity)]
+    /// use std::mem::MaybeUninit;
+    ///
+    /// let mut vec = Vec::with_capacity(32);
+    /// let src = [0; 16];
+    ///
+    /// MaybeUninit::write_slice(&mut vec.spare_capacity_mut()[..src.len()], &src);
+    ///
+    /// // SAFETY: we have just copied all the elements of len into the spare capacity
+    /// // the first src.len() elements of the vec are valid now.
+    /// unsafe {
+    ///     vec.set_len(src.len());
+    /// }
+    ///
+    /// assert_eq!(vec, src);
+    /// ```
+    ///
+    /// [`write_slice_cloned`]: MaybeUninit::write_slice_cloned
+    /// [`slice::copy_from_slice`]: ../../std/primitive.slice.html#method.copy_from_slice
+    #[unstable(feature = "maybe_uninit_write_slice", issue = "79995")]
+    pub fn write_slice<'a>(this: &'a mut [MaybeUninit<T>], src: &[T]) -> &'a mut [T]
+    where
+        T: Copy,
+    {
+        // SAFETY: &[T] and &[MaybeUninit<T>] have the same layout
+        let uninit_src: &[MaybeUninit<T>] = unsafe { super::transmute(src) };
+
+        this.copy_from_slice(uninit_src);
+
+        // SAFETY: Valid elements have just been copied into `this` so it is initalized
+        unsafe { MaybeUninit::slice_assume_init_mut(this) }
+    }
+
+    /// Clones the elements from `src` to `this`, returning a mutable reference to the now initalized contents of `this`.
+    /// Any already initalized elements will not be dropped.
+    ///
+    /// If `T` implements `Copy`, use [`write_slice`]
+    ///
+    /// This is similar to [`slice::clone_from_slice`] but does not drop existing elements.
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if the two slices have different lengths, or if the implementation of `Clone` panics.
+    ///
+    /// If there is a panic, the already cloned elements will be dropped.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(maybe_uninit_write_slice)]
+    /// use std::mem::MaybeUninit;
+    ///
+    /// let mut dst = [MaybeUninit::uninit(), MaybeUninit::uninit(), MaybeUninit::uninit(), MaybeUninit::uninit(), MaybeUninit::uninit()];
+    /// let src = ["wibbly".to_string(), "wobbly".to_string(), "timey".to_string(), "wimey".to_string(), "stuff".to_string()];
+    ///
+    /// let init = MaybeUninit::write_slice_cloned(&mut dst, &src);
+    ///
+    /// assert_eq!(init, src);
+    /// ```
+    ///
+    /// ```
+    /// #![feature(maybe_uninit_write_slice, vec_spare_capacity)]
+    /// use std::mem::MaybeUninit;
+    ///
+    /// let mut vec = Vec::with_capacity(32);
+    /// let src = ["rust", "is", "a", "pretty", "cool", "language"];
+    ///
+    /// MaybeUninit::write_slice_cloned(&mut vec.spare_capacity_mut()[..src.len()], &src);
+    ///
+    /// // SAFETY: we have just cloned all the elements of len into the spare capacity
+    /// // the first src.len() elements of the vec are valid now.
+    /// unsafe {
+    ///     vec.set_len(src.len());
+    /// }
+    ///
+    /// assert_eq!(vec, src);
+    /// ```
+    ///
+    /// [`write_slice`]: MaybeUninit::write_slice
+    /// [`slice::clone_from_slice`]: ../../std/primitive.slice.html#method.clone_from_slice
+    #[unstable(feature = "maybe_uninit_write_slice", issue = "79995")]
+    pub fn write_slice_cloned<'a>(this: &'a mut [MaybeUninit<T>], src: &[T]) -> &'a mut [T]
+    where
+        T: Clone,
+    {
+        // unlike copy_from_slice this does not call clone_from_slice on the slice
+        // this is because `MaybeUninit<T: Clone>` does not implement Clone.
+
+        struct Guard<'a, T> {
+            slice: &'a mut [MaybeUninit<T>],
+            initialized: usize,
+        }
+
+        impl<'a, T> Drop for Guard<'a, T> {
+            fn drop(&mut self) {
+                let initialized_part = &mut self.slice[..self.initialized];
+                // SAFETY: this raw slice will contain only initialized objects
+                // that's why, it is allowed to drop it.
+                unsafe {
+                    crate::ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(initialized_part));
+                }
+            }
+        }
+
+        assert_eq!(this.len(), src.len(), "destination and source slices have different lengths");
+        // NOTE: We need to explicitly slice them to the same length
+        // for bounds checking to be elided, and the optimizer will
+        // generate memcpy for simple cases (for example T = u8).
+        let len = this.len();
+        let src = &src[..len];
+
+        // guard is needed b/c panic might happen during a clone
+        let mut guard = Guard { slice: this, initialized: 0 };
+
+        for i in 0..len {
+            guard.slice[i].write(src[i].clone());
+            guard.initialized += 1;
+        }
+
+        super::forget(guard);
+
+        // SAFETY: Valid elements have just been written into `this` so it is initalized
+        unsafe { MaybeUninit::slice_assume_init_mut(this) }
+    }
 }
diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs
index 2aa3598a0d9..2828235c3e3 100644
--- a/library/core/tests/lib.rs
+++ b/library/core/tests/lib.rs
@@ -8,6 +8,7 @@
 #![feature(bound_cloned)]
 #![feature(box_syntax)]
 #![feature(cell_update)]
+#![feature(cfg_panic)]
 #![feature(cfg_target_has_atomic)]
 #![feature(const_assume)]
 #![feature(const_cell_into_inner)]
@@ -33,6 +34,7 @@
 #![feature(raw)]
 #![feature(sort_internals)]
 #![feature(slice_partition_at_index)]
+#![feature(maybe_uninit_write_slice)]
 #![feature(min_specialization)]
 #![feature(step_trait)]
 #![feature(step_trait_ext)]
diff --git a/library/core/tests/mem.rs b/library/core/tests/mem.rs
index 268c2ed283f..5e24fa690ef 100644
--- a/library/core/tests/mem.rs
+++ b/library/core/tests/mem.rs
@@ -1,5 +1,7 @@
 use core::mem::*;
 
+use std::rc::Rc;
+
 #[test]
 fn size_of_basic() {
     assert_eq!(size_of::<u8>(), 1);
@@ -137,3 +139,125 @@ fn assume_init_good() {
 
     assert!(TRUE);
 }
+
+#[test]
+fn uninit_write_slice() {
+    let mut dst = [MaybeUninit::new(255); 64];
+    let src = [0; 64];
+
+    assert_eq!(MaybeUninit::write_slice(&mut dst, &src), &src);
+}
+
+#[test]
+#[should_panic(expected = "source slice length (32) does not match destination slice length (64)")]
+fn uninit_write_slice_panic_lt() {
+    let mut dst = [MaybeUninit::uninit(); 64];
+    let src = [0; 32];
+
+    MaybeUninit::write_slice(&mut dst, &src);
+}
+
+#[test]
+#[should_panic(expected = "source slice length (128) does not match destination slice length (64)")]
+fn uninit_write_slice_panic_gt() {
+    let mut dst = [MaybeUninit::uninit(); 64];
+    let src = [0; 128];
+
+    MaybeUninit::write_slice(&mut dst, &src);
+}
+
+#[test]
+fn uninit_clone_from_slice() {
+    let mut dst = [MaybeUninit::new(255); 64];
+    let src = [0; 64];
+
+    assert_eq!(MaybeUninit::write_slice_cloned(&mut dst, &src), &src);
+}
+
+#[test]
+#[should_panic(expected = "destination and source slices have different lengths")]
+fn uninit_write_slice_cloned_panic_lt() {
+    let mut dst = [MaybeUninit::uninit(); 64];
+    let src = [0; 32];
+
+    MaybeUninit::write_slice_cloned(&mut dst, &src);
+}
+
+#[test]
+#[should_panic(expected = "destination and source slices have different lengths")]
+fn uninit_write_slice_cloned_panic_gt() {
+    let mut dst = [MaybeUninit::uninit(); 64];
+    let src = [0; 128];
+
+    MaybeUninit::write_slice_cloned(&mut dst, &src);
+}
+
+#[test]
+#[cfg(panic = "unwind")]
+fn uninit_write_slice_cloned_mid_panic() {
+    use std::panic;
+
+    enum IncrementOrPanic {
+        Increment(Rc<()>),
+        ExpectedPanic,
+        UnexpectedPanic,
+    }
+
+    impl Clone for IncrementOrPanic {
+        fn clone(&self) -> Self {
+            match self {
+                Self::Increment(rc) => Self::Increment(rc.clone()),
+                Self::ExpectedPanic => panic!("expected panic on clone"),
+                Self::UnexpectedPanic => panic!("unexpected panic on clone"),
+            }
+        }
+    }
+
+    let rc = Rc::new(());
+
+    let mut dst = [
+        MaybeUninit::uninit(),
+        MaybeUninit::uninit(),
+        MaybeUninit::uninit(),
+        MaybeUninit::uninit(),
+    ];
+
+    let src = [
+        IncrementOrPanic::Increment(rc.clone()),
+        IncrementOrPanic::Increment(rc.clone()),
+        IncrementOrPanic::ExpectedPanic,
+        IncrementOrPanic::UnexpectedPanic,
+    ];
+
+    let err = panic::catch_unwind(panic::AssertUnwindSafe(|| {
+        MaybeUninit::write_slice_cloned(&mut dst, &src);
+    }));
+
+    drop(src);
+
+    match err {
+        Ok(_) => unreachable!(),
+        Err(payload) => {
+            payload
+                .downcast::<&'static str>()
+                .and_then(|s| if *s == "expected panic on clone" { Ok(s) } else { Err(s) })
+                .unwrap_or_else(|p| panic::resume_unwind(p));
+
+            assert_eq!(Rc::strong_count(&rc), 1)
+        }
+    }
+}
+
+#[test]
+fn uninit_write_slice_cloned_no_drop() {
+    let rc = Rc::new(());
+
+    let mut dst = [MaybeUninit::uninit()];
+    let src = [rc.clone()];
+
+    MaybeUninit::write_slice_cloned(&mut dst, &src);
+
+    drop(src);
+
+    assert_eq!(Rc::strong_count(&rc), 2);
+}