Implement MaybeUninit::fill{,_cloned,_mut,_with,_from}

ACP: rust-lang/libs-team#156

Author: jmillikin

library/core/src/mem/maybe_uninit.rs

@@ -1123,22 +1123,6 @@ impl<T> MaybeUninit<T> {
         // 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
@@ -1147,7 +1131,7 @@ impl<T> MaybeUninit<T> {
         let src = &src[..len];
 
         // guard is needed b/c panic might happen during a clone
-        let mut guard = Guard { slice: this, initialized: 0 };
+        let mut guard = CloneGuard { slice: this, initialized: 0 };
 
         for i in 0..len {
             guard.slice[i].write(src[i].clone());
@@ -1160,6 +1144,202 @@ impl<T> MaybeUninit<T> {
         unsafe { MaybeUninit::slice_assume_init_mut(this) }
     }
 
+    /// Fills `this` with elements by copying `value`, returning a reference to
+    /// the now initialized contents of `this`.
+    ///
+    /// This is similar to [`slice::fill`] but is restricted to `Copy` values.
+    /// Use [`MaybeUninit::fill_cloned`] to initialize from a `Clone` value.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(maybe_uninit_fill)]
+    /// use std::mem::MaybeUninit;
+    ///
+    /// let mut dst = [MaybeUninit::uninit(); 5];
+    /// let init = MaybeUninit::fill(&mut dst, 0u8);
+    ///
+    /// assert_eq!(init, &[0, 0, 0, 0, 0]);
+    /// ```
+    #[doc(alias = "memset")]
+    #[unstable(feature = "maybe_uninit_fill", issue = "117428")]
+    pub fn fill<'a>(this: &'a mut [MaybeUninit<T>], value: T) -> &'a [T]
+    where
+        T: Copy,
+    {
+        this.fill(MaybeUninit::new(value));
+        // SAFETY: Valid elements have just been copied into `this` so it is initialized
+        unsafe { MaybeUninit::slice_assume_init_ref(this) }
+    }
+
+    /// Fills `this` with elements by copying `value`, returning a mutable
+    /// reference to the now initialized contents of `this`.
+    ///
+    /// This is similar to [`slice::fill`] but is restricted to `Copy` values.
+    /// Use [`MaybeUninit::fill_cloned`] to initialize from a `Clone` value.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(maybe_uninit_fill)]
+    /// use std::mem::MaybeUninit;
+    ///
+    /// let mut dst = [MaybeUninit::uninit(); 5];
+    /// let init = MaybeUninit::fill_mut(&mut dst, 0u8);
+    /// init[4] = 123;
+    ///
+    /// assert_eq!(init, &[0, 0, 0, 0, 123]);
+    /// ```
+    #[unstable(feature = "maybe_uninit_fill", issue = "117428")]
+    pub fn fill_mut<'a>(this: &'a mut [MaybeUninit<T>], value: T) -> &'a mut [T]
+    where
+        T: Copy,
+    {
+        this.fill(MaybeUninit::new(value));
+        // SAFETY: Valid elements have just been copied into `this` so it is initialized
+        unsafe { MaybeUninit::slice_assume_init_mut(this) }
+    }
+
+    /// Fills `this` with elements by cloning `value`, returning a reference to
+    /// the now initialized contents of `this`. Any already  initialized elements
+    /// will not be dropped.
+    ///
+    /// This is similar to [`slice::fill`] but does not drop existing elements.
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if the implementation of `Clone` panics.
+    ///
+    /// If there is a panic, the already initialized elements will be dropped.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(maybe_uninit_fill)]
+    /// use std::mem::MaybeUninit;
+    ///
+    /// let mut dst = [
+    ///    MaybeUninit::uninit(),
+    ///    MaybeUninit::uninit(),
+    ///    MaybeUninit::uninit(),
+    /// ];
+    /// let msg = String::from("hello");
+    /// let init = MaybeUninit::fill_cloned(&mut dst, &msg);
+    ///
+    /// assert_eq!(init, &["hello", "hello", "hello"]);
+    /// ```
+    #[unstable(feature = "maybe_uninit_fill", issue = "117428")]
+    pub fn fill_cloned<'a>(this: &'a mut [MaybeUninit<T>], value: &T) -> &'a [T]
+    where
+        T: Clone,
+    {
+        let len = this.len();
+
+        // guard is needed b/c panic might happen during a clone
+        let mut guard = CloneGuard { slice: this, initialized: 0 };
+
+        for i in 0..len {
+            guard.slice[i].write(value.clone());
+            guard.initialized += 1;
+        }
+
+        super::forget(guard);
+
+        // SAFETY: Valid elements have just been written into `this` so it is initialized
+        unsafe { MaybeUninit::slice_assume_init_ref(this) }
+    }
+
+    /// Fills `this` with elements returned by calling a closure repeatedly,
+    /// returning a reference to the now initialized contents of `this`. Any
+    /// already initialized elements will not be dropped.
+    ///
+    /// This is similar to [`slice::fill_with`] but does not drop existing
+    /// elements.
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if the closure panics.
+    ///
+    /// If there is a panic, the already initialized elements will be dropped.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(maybe_uninit_fill)]
+    /// use std::mem::MaybeUninit;
+    ///
+    /// let mut dst = [MaybeUninit::uninit(); 5];
+    /// let mut next = 0;
+    /// let init = MaybeUninit::fill_with(&mut dst, || { next += 1; next });
+    ///
+    /// assert_eq!(init, &[1, 2, 3, 4, 5]);
+    /// ```
+    #[unstable(feature = "maybe_uninit_fill", issue = "117428")]
+    pub fn fill_with<'a, F>(this: &'a mut [MaybeUninit<T>], mut f: F) -> &'a [T]
+    where
+        F: FnMut() -> T,
+    {
+        let len = this.len();
+        let mut guard = CloneGuard { slice: this, initialized: 0 };
+
+        for i in 0..len {
+            guard.slice[i].write(f());
+            guard.initialized += 1;
+        }
+
+        super::forget(guard);
+
+        // SAFETY: Valid elements have just been written into `this` so it is initialized
+        unsafe { MaybeUninit::slice_assume_init_ref(this) }
+    }
+
+    /// Fills `this` with the contents of an iterator, returning a reference to
+    /// the now initialized contents of `this`. Any already  initialized elements
+    /// will not be dropped.
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if the iterator panics.
+    ///
+    /// If there is a panic, the already initialized elements will be dropped.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(maybe_uninit_fill)]
+    /// use std::mem::MaybeUninit;
+    ///
+    /// let mut dst = [MaybeUninit::uninit(); 5];
+    /// let mut iter = [1, 2, 3].into_iter().cycle();
+    /// let init = MaybeUninit::fill_from(&mut dst, iter);
+    ///
+    /// assert_eq!(init, &[1, 2, 3, 1, 2]);
+    /// ```
+    #[unstable(feature = "maybe_uninit_fill", issue = "117428")]
+    pub fn fill_from<'a, I>(this: &'a mut [MaybeUninit<T>], mut iter: I) -> &'a [T]
+    where
+        I: Iterator<Item = T>,
+    {
+        let len = this.len();
+        let mut guard = CloneGuard { slice: this, initialized: 0 };
+
+        for i in 0..len {
+            match iter.next() {
+                Some(value) => {
+                    guard.slice[i].write(value);
+                    guard.initialized += 1;
+                }
+                None => break,
+            }
+        }
+
+        let init_len = guard.initialized;
+        super::forget(guard);
+
+        // SAFETY: Valid elements have just been written into `this` so it is initialized
+        unsafe { MaybeUninit::slice_assume_init_ref(&mut this[..init_len]) }
+    }
+
     /// Returns the contents of this `MaybeUninit` as a slice of potentially uninitialized bytes.
     ///
     /// Note that even if the contents of a `MaybeUninit` have been initialized, the value may still
@@ -1313,3 +1493,19 @@ impl<T, const N: usize> [MaybeUninit<T>; N] {
         unsafe { intrinsics::transmute_unchecked(self) }
     }
 }
+
+struct CloneGuard<'a, T> {
+    slice: &'a mut [MaybeUninit<T>],
+    initialized: usize,
+}
+
+impl<'a, T> Drop for CloneGuard<'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));
+        }
+    }
+}

library/core/tests/lib.rs

@@ -49,6 +49,7 @@
 #![feature(slice_from_ptr_range)]
 #![feature(slice_split_once)]
 #![feature(split_as_slice)]
+#![feature(maybe_uninit_fill)]
 #![feature(maybe_uninit_uninit_array)]
 #![feature(maybe_uninit_write_slice)]
 #![feature(maybe_uninit_uninit_array_transpose)]

library/core/tests/mem.rs

@@ -308,21 +308,226 @@ fn uninit_write_slice_cloned_mid_panic() {
     }
 }
 
+#[derive(Clone)]
+struct DropBomb;
+
+impl Drop for DropBomb {
+    fn drop(&mut self) {
+        panic!("dropped a bomb! kaboom")
+    }
+}
+
 #[test]
 fn uninit_write_slice_cloned_no_drop() {
-    #[derive(Clone)]
-    struct Bomb;
+    let mut dst = [MaybeUninit::uninit()];
+    let src = [DropBomb];
+
+    MaybeUninit::write_slice_cloned(&mut dst, &src);
+
+    forget(src);
+}
+
+#[test]
+fn uninit_fill() {
+    let mut dst = [MaybeUninit::new(255); 64];
+    let expect = [0; 64];
+
+    assert_eq!(MaybeUninit::fill(&mut dst, 0), &expect);
+}
+
+#[test]
+fn uninit_fill_mut() {
+    let mut dst = [MaybeUninit::new(255); 64];
+    let expect = [0; 64];
+
+    let init: &mut [u8] = MaybeUninit::fill_mut(&mut dst, 0);
+    assert_eq!(init, &expect);
+}
 
-    impl Drop for Bomb {
-        fn drop(&mut self) {
-            panic!("dropped a bomb! kaboom")
+#[test]
+fn uninit_fill_cloned() {
+    let mut dst = [MaybeUninit::new(255); 64];
+    let expect = [0; 64];
+
+    assert_eq!(MaybeUninit::fill_cloned(&mut dst, &0), &expect);
+}
+
+struct IncrementUntilPanic {
+    limit: usize,
+    rc: Rc<()>,
+}
+
+impl Clone for IncrementUntilPanic {
+    fn clone(&self) -> Self {
+        if Rc::strong_count(&self.rc) >= self.limit {
+            panic!("expected panic on clone");
+        }
+        Self {
+            limit: self.limit,
+            rc: self.rc.clone(),
         }
     }
+}
 
+#[test]
+#[cfg(panic = "unwind")]
+fn uninit_fill_cloned_mid_panic() {
+    use std::panic;
+
+    let rc = Rc::new(());
+
+    let mut dst = [
+        MaybeUninit::uninit(),
+        MaybeUninit::uninit(),
+        MaybeUninit::uninit(),
+    ];
+
+    let src = IncrementUntilPanic { limit: 3, rc: rc.clone() };
+    let err = panic::catch_unwind(panic::AssertUnwindSafe(|| {
+        MaybeUninit::fill_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]
+#[cfg(panic = "unwind")]
+fn uninit_fill_cloned_no_drop() {
     let mut dst = [MaybeUninit::uninit()];
-    let src = [Bomb];
+    let src = DropBomb;
 
-    MaybeUninit::write_slice_cloned(&mut dst, &src);
+    MaybeUninit::fill_cloned(&mut dst, &src);
+
+    forget(src);
+}
+
+#[test]
+fn uninit_fill_with() {
+    let mut dst = [MaybeUninit::new(255); 64];
+    let expect = [0; 64];
+
+    assert_eq!(MaybeUninit::fill_with(&mut dst, || 0), &expect);
+}
+
+#[test]
+#[cfg(panic = "unwind")]
+fn uninit_fill_with_mid_panic() {
+    use std::panic;
+
+    let rc = Rc::new(());
+
+    let mut dst = [
+        MaybeUninit::uninit(),
+        MaybeUninit::uninit(),
+        MaybeUninit::uninit(),
+    ];
+
+    let src = IncrementUntilPanic { limit: 3, rc: rc.clone() };
+    let err = panic::catch_unwind(panic::AssertUnwindSafe(|| {
+        MaybeUninit::fill_with(&mut dst, || src.clone());
+    }));
+
+    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]
+#[cfg(panic = "unwind")]
+fn uninit_fill_with_no_drop() {
+    let mut dst = [MaybeUninit::uninit()];
+    let src = DropBomb;
+
+    MaybeUninit::fill_with(&mut dst, || src.clone());
+
+    forget(src);
+}
+
+#[test]
+fn uninit_fill_from() {
+    let mut dst = [MaybeUninit::new(255); 64];
+    let src = [0; 64];
+
+    assert_eq!(MaybeUninit::fill_from(&mut dst, src.into_iter()), &src);
+}
+
+#[test]
+#[cfg(panic = "unwind")]
+fn uninit_fill_from_mid_panic() {
+    use std::panic;
+
+    struct IterUntilPanic {
+        limit: usize,
+        rc: Rc<()>,
+    }
+
+    impl Iterator for IterUntilPanic {
+        type Item = Rc<()>;
+        fn next(&mut self) -> Option<Self::Item> {
+            if Rc::strong_count(&self.rc) >= self.limit {
+                panic!("expected panic on next");
+            }
+            Some(self.rc.clone())
+        }
+    }
+
+    let rc = Rc::new(());
+
+    let mut dst = [
+        MaybeUninit::uninit(),
+        MaybeUninit::uninit(),
+        MaybeUninit::uninit(),
+        MaybeUninit::uninit(),
+    ];
+
+    let src = IterUntilPanic { limit: 3, rc: rc.clone() };
+
+    let err = panic::catch_unwind(panic::AssertUnwindSafe(|| {
+        MaybeUninit::fill_from(&mut dst, src);
+    }));
+
+    match err {
+        Ok(_) => unreachable!(),
+        Err(payload) => {
+            payload
+                .downcast::<&'static str>()
+                .and_then(|s| if *s == "expected panic on next" { Ok(s) } else { Err(s) })
+                .unwrap_or_else(|p| panic::resume_unwind(p));
+
+            assert_eq!(Rc::strong_count(&rc), 1)
+        }
+    }
+}
+
+#[test]
+#[cfg(panic = "unwind")]
+fn uninit_fill_from_no_drop() {
+    let mut dst = [MaybeUninit::uninit()];
+    let src = [DropBomb];
+
+    MaybeUninit::fill_from(&mut dst, src.iter());
 
     forget(src);
 }