Add core::clone::CloneToUninit.

kpreid · kpreid · commit ec201b8650e5 · 2024-06-22T08:08:00.000-07:00
This trait allows cloning DSTs, but is unsafe to implement and use
because it writes to possibly-uninitialized memory which must be of the
correct size, and must initialize that memory.

It is only implemented for `T: Clone` and `[T] where T: Clone`, but
additional implementations could be provided for specific `dyn Trait`
or custom-DST types.
diff --git a/library/core/src/clone.rs b/library/core/src/clone.rs
@@ -36,6 +36,9 @@
 
 #![stable(feature = "rust1", since = "1.0.0")]
 
+use crate::mem::{self, MaybeUninit};
+use crate::ptr;
+
 /// A common trait for the ability to explicitly duplicate an object.
 ///
 /// Differs from [`Copy`] in that [`Copy`] is implicit and an inexpensive bit-wise copy, while
@@ -204,6 +207,189 @@ pub struct AssertParamIsCopy<T: Copy + ?Sized> {
     _field: crate::marker::PhantomData<T>,
 }
 
+/// A generalization of [`Clone`] to dynamically-sized types stored in arbitrary containers.
+///
+/// This trait is implemented for all types implementing [`Clone`], and also [slices](slice) of all
+/// such types. You may also implement this trait to enable cloning trait objects and custom DSTs
+/// (structures containing dynamically-sized fields).
+///
+/// # Safety
+///
+/// Implementations must ensure that when `.clone_to_uninit(dst)` returns normally rather than
+/// panicking, it always leaves `*dst` initialized as a valid value of type `Self`.
+///
+/// # See also
+///
+/// * [`Clone::clone_from`] is a safe function which may be used instead when `Self` is a [`Sized`]
+///   and the destination is already initialized; it may be able to reuse allocations owned by
+///   the destination.
+/// * [`ToOwned`], which allocates a new destination container.
+///
+/// [`ToOwned`]: ../../std/borrow/trait.ToOwned.html
+#[unstable(feature = "clone_to_uninit", issue = "126799")]
+pub unsafe trait CloneToUninit {
+    /// Performs copy-assignment from `self` to `dst`.
+    ///
+    /// This is analogous to to `std::ptr::write(dst, self.clone())`,
+    /// except that `self` may be a dynamically-sized type ([`!Sized`](Sized)).
+    ///
+    /// Before this function is called, `dst` may point to uninitialized memory.
+    /// After this function is called, `dst` will point to initialized memory; it will be
+    /// sound to create a `&Self` reference from the pointer.
+    ///
+    /// # Safety
+    ///
+    /// Behavior is undefined if any of the following conditions are violated:
+    ///
+    /// * `dst` must be [valid] for writes.
+    /// * `dst` must be properly aligned.
+    /// * `dst` must have the same [pointer metadata] (slice length or `dyn` vtable) as `self`.
+    ///
+    /// [valid]: ptr#safety
+    /// [pointer metadata]: crate::ptr::metadata()
+    ///
+    /// # Panics
+    ///
+    /// This function may panic. (For example, it might panic if memory allocation for a clone
+    /// of a value owned by `self` fails.)
+    /// If the call panics, then `*dst` should be treated as uninitialized memory; it must not be
+    /// read or dropped, because even if it was previously valid, it may have been partially
+    /// overwritten.
+    ///
+    /// The caller may also need to take care to deallocate the allocation pointed to by `dst`,
+    /// if applicable, to avoid a memory leak, and may need to take other precautions to ensure
+    /// soundness in the presence of unwinding.
+    ///
+    /// Implementors should avoid leaking values by, upon unwinding, dropping all component values
+    /// that might have already been created. (For example, if a `[Foo]` of length 3 is being
+    /// cloned, and the second of the three calls to `Foo::clone()` unwinds, then the first `Foo`
+    /// cloned should be dropped.)
+    unsafe fn clone_to_uninit(&self, dst: *mut Self);
+}
+
+#[unstable(feature = "clone_to_uninit", issue = "126799")]
+unsafe impl<T: Clone> CloneToUninit for T {
+    default unsafe fn clone_to_uninit(&self, dst: *mut Self) {
+        // SAFETY: The safety conditions of clone_to_uninit() are a superset of those of
+        // ptr::write().
+        unsafe {
+            // We hope the optimizer will figure out to create the cloned value in-place,
+            // skipping ever storing it on the stack and the copy to the destination.
+            ptr::write(dst, self.clone());
+        }
+    }
+}
+
+// Specialized implementation for types that are [`Copy`], not just [`Clone`],
+// and can therefore be copied bitwise.
+#[unstable(feature = "clone_to_uninit", issue = "126799")]
+unsafe impl<T: Copy> CloneToUninit for T {
+    unsafe fn clone_to_uninit(&self, dst: *mut Self) {
+        // SAFETY: The safety conditions of clone_to_uninit() are a superset of those of
+        // ptr::copy_nonoverlapping().
+        unsafe {
+            ptr::copy_nonoverlapping(self, dst, 1);
+        }
+    }
+}
+
+#[unstable(feature = "clone_to_uninit", issue = "126799")]
+unsafe impl<T: Clone> CloneToUninit for [T] {
+    #[cfg_attr(debug_assertions, track_caller)]
+    default unsafe fn clone_to_uninit(&self, dst: *mut Self) {
+        let len = self.len();
+        // This is the most likely mistake to make, so check it as a debug assertion.
+        debug_assert_eq!(
+            len,
+            dst.len(),
+            "clone_to_uninit() source and destination must have equal lengths",
+        );
+
+        // SAFETY: The produced `&mut` is valid because:
+        // * The caller is obligated to provide a pointer which is valid for writes.
+        // * All bytes pointed to are in MaybeUninit, so we don't care about the memory's
+        //   initialization status.
+        let uninit_ref = unsafe { &mut *(dst as *mut [MaybeUninit<T>]) };
+
+        // Copy the elements
+        let mut initializing = InitializingSlice::from_fully_uninit(uninit_ref);
+        for element_ref in self.iter() {
+            // If the clone() panics, `initializing` will take care of the cleanup.
+            initializing.push(element_ref.clone());
+        }
+        // If we reach here, then the entire slice is initialized, and we've satisfied our
+        // responsibilities to the caller. Disarm the cleanup guard by forgetting it.
+        mem::forget(initializing);
+    }
+}
+
+#[unstable(feature = "clone_to_uninit", issue = "126799")]
+unsafe impl<T: Copy> CloneToUninit for [T] {
+    #[cfg_attr(debug_assertions, track_caller)]
+    unsafe fn clone_to_uninit(&self, dst: *mut Self) {
+        let len = self.len();
+        // This is the most likely mistake to make, so check it as a debug assertion.
+        debug_assert_eq!(
+            len,
+            dst.len(),
+            "clone_to_uninit() source and destination must have equal lengths",
+        );
+
+        // SAFETY: The safety conditions of clone_to_uninit() are a superset of those of
+        // ptr::copy_nonoverlapping().
+        unsafe {
+            ptr::copy_nonoverlapping(self.as_ptr(), dst.as_mut_ptr(), len);
+        }
+    }
+}
+
+/// Ownership of a collection of values stored in a non-owned `[MaybeUninit<T>]`, some of which
+/// are not yet initialized. This is sort of like a `Vec` that doesn't own its allocation.
+/// Its responsibility is to provide cleanup on unwind by dropping the values that *are*
+/// initialized, unless disarmed by forgetting.
+///
+/// This is a helper for `impl<T: Clone> CloneToUninit for [T]`.
+struct InitializingSlice<'a, T> {
+    data: &'a mut [MaybeUninit<T>],
+    /// Number of elements of `*self.data` that are initialized.
+    initialized_len: usize,
+}
+
+impl<'a, T> InitializingSlice<'a, T> {
+    #[inline]
+    fn from_fully_uninit(data: &'a mut [MaybeUninit<T>]) -> Self {
+        Self { data, initialized_len: 0 }
+    }
+
+    /// Push a value onto the end of the initialized part of the slice.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the slice is already fully initialized.
+    #[inline]
+    fn push(&mut self, value: T) {
+        MaybeUninit::write(&mut self.data[self.initialized_len], value);
+        self.initialized_len += 1;
+    }
+}
+
+impl<'a, T> Drop for InitializingSlice<'a, T> {
+    #[cold] // will only be invoked on unwind
+    fn drop(&mut self) {
+        let initialized_slice = ptr::slice_from_raw_parts_mut(
+            MaybeUninit::slice_as_mut_ptr(self.data),
+            self.initialized_len,
+        );
+        // SAFETY:
+        // * the pointer is valid because it was made from a mutable reference
+        // * `initialized_len` counts the initialized elements as an invariant of this type,
+        //   so each of the pointed-to elements is initialized and may be dropped.
+        unsafe {
+            ptr::drop_in_place::<[T]>(initialized_slice);
+        }
+    }
+}
+
 /// Implementations of `Clone` for primitive types.
 ///
 /// Implementations that cannot be described in Rust
diff --git a/library/core/tests/clone.rs b/library/core/tests/clone.rs
@@ -1,3 +1,6 @@
+use core::clone::CloneToUninit;
+use core::mem::MaybeUninit;
+
 #[test]
 #[allow(suspicious_double_ref_op)]
 fn test_borrowed_clone() {
@@ -14,3 +17,66 @@ fn test_clone_from() {
     b.clone_from(&a);
     assert_eq!(*b, 5);
 }
+
+#[test]
+fn test_clone_to_uninit_slice_success() {
+    // Using `String`s to exercise allocation and Drop of the individual elements;
+    // if something is aliased or double-freed, at least Miri will catch that.
+    let a: [String; 3] = ["a", "b", "c"].map(String::from);
+
+    let mut storage: MaybeUninit<[String; 3]> = MaybeUninit::uninit();
+    let b: [String; 3] = unsafe {
+        a[..].clone_to_uninit(storage.as_mut_ptr() as *mut [String]);
+        storage.assume_init()
+    };
+
+    assert_eq!(a, b);
+}
+
+#[test]
+#[cfg(panic = "unwind")]
+fn test_clone_to_uninit_slice_drops_on_panic() {
+    use core::sync::atomic::{AtomicUsize, Ordering::Relaxed};
+
+    /// A static counter is OK to use as long as _this one test_ isn't run several times in
+    /// multiple threads.
+    static COUNTER: AtomicUsize = AtomicUsize::new(0);
+    /// Counts how many instances are live, and panics if a fifth one is created
+    struct CountsDropsAndPanics {}
+    impl CountsDropsAndPanics {
+        fn new() -> Self {
+            COUNTER.fetch_add(1, Relaxed);
+            Self {}
+        }
+    }
+    impl Clone for CountsDropsAndPanics {
+        fn clone(&self) -> Self {
+            if COUNTER.load(Relaxed) == 4 { panic!("intentional panic") } else { Self::new() }
+        }
+    }
+    impl Drop for CountsDropsAndPanics {
+        fn drop(&mut self) {
+            COUNTER.fetch_sub(1, Relaxed);
+        }
+    }
+
+    let a: [CountsDropsAndPanics; 3] = core::array::from_fn(|_| CountsDropsAndPanics::new());
+    assert_eq!(COUNTER.load(Relaxed), 3);
+
+    let panic_payload = std::panic::catch_unwind(|| {
+        let mut storage: MaybeUninit<[CountsDropsAndPanics; 3]> = MaybeUninit::uninit();
+        // This should panic halfway through
+        unsafe {
+            a[..].clone_to_uninit(storage.as_mut_ptr() as *mut [CountsDropsAndPanics]);
+        }
+    })
+    .unwrap_err();
+    assert_eq!(panic_payload.downcast().unwrap(), Box::new("intentional panic"));
+
+    // Check for lack of leak, which is what this test is looking for
+    assert_eq!(COUNTER.load(Relaxed), 3, "leaked during clone!");
+
+    // Might as well exercise the rest of the drops
+    drop(a);
+    assert_eq!(COUNTER.load(Relaxed), 0);
+}
diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs
@@ -8,6 +8,7 @@
 #![feature(async_iterator)]
 #![feature(bigint_helper_methods)]
 #![feature(cell_update)]
+#![feature(clone_to_uninit)]
 #![feature(const_align_offset)]
 #![feature(const_align_of_val_raw)]
 #![feature(const_black_box)]
@@ -54,6 +55,7 @@
 #![feature(slice_split_once)]
 #![feature(split_as_slice)]
 #![feature(maybe_uninit_fill)]
+#![feature(maybe_uninit_slice)]
 #![feature(maybe_uninit_uninit_array)]
 #![feature(maybe_uninit_write_slice)]
 #![feature(maybe_uninit_uninit_array_transpose)]
