core: generalize BorrowedCursor::ensure_init

library/core/src/io/borrowed_buf.rs

@@ -2,7 +2,6 @@
 
 use crate::fmt::{self, Debug, Formatter};
 use crate::mem::{self, MaybeUninit};
-use crate::ptr;
 
 /// A borrowed buffer of initially uninitialized elements, which is incrementally filled.
 ///
@@ -357,24 +356,21 @@ impl<'a, T: Copy> BorrowedCursor<'a, T> {
     }
 }
 
-impl<'a> BorrowedCursor<'a, u8> {
-    /// Initializes all bytes in the cursor and returns them.
+impl<'a, T: Default + Copy> BorrowedCursor<'a, T> {
+    /// Initializes all elements in the cursor with their default value and
+    /// returns them.
     #[unstable(feature = "borrowed_buf_init", issue = "160476")]
     #[inline]
-    pub fn ensure_init(&mut self) -> &mut [u8] {
-        // SAFETY: always in bounds and we never uninitialize these bytes.
+    pub fn ensure_init(&mut self) -> &mut [T] {
+        // SAFETY: always in bounds and we never uninitialize these elements.
         let unfilled = unsafe { self.buf.buf.get_unchecked_mut(self.buf.filled..) };
 
         if !self.buf.init {
-            // SAFETY: 0 is a valid value for MaybeUninit<u8> and the length matches the allocation
-            // since it is comes from a slice reference.
-            unsafe {
-                ptr::write_bytes(unfilled.as_mut_ptr(), 0, unfilled.len());
-            }
+            unfilled.write_default();
             self.buf.init = true;
         }
 
-        // SAFETY: these bytes have just been initialized if they weren't before
+        // SAFETY: these elements have just been initialized if they weren't before
         unsafe { unfilled.assume_init_mut() }
     }
 }

library/core/src/mem/maybe_uninit.rs

@@ -1286,8 +1286,8 @@ impl<T> [MaybeUninit<T>] {
     /// Fills a slice with elements returned by calling a closure for each index.
     ///
     /// This method uses a closure to create new values. If you'd rather `Clone` a given value, use
-    /// [slice::write_filled]. If you want to use the `Default` trait to generate values, you can
-    /// pass [`|_| Default::default()`][Default::default] as the argument.
+    /// [`slice::write_filled`]. If you want to use the `Default` trait to generate values, use
+    /// [`slice::write_default`].
     ///
     /// # Panics
     ///
@@ -1324,6 +1324,73 @@ impl<T> [MaybeUninit<T>] {
         unsafe { self.assume_init_mut() }
     }
 
+    /// Fills a slice with elements returned by calling [`Default::default`] for each index.
+    ///
+    /// # Panics
+    ///
+    /// This function will panic if any call to [`Default::default`] panics.
+    ///
+    /// If such a panic occurs, any elements previously initialized during this operation will be
+    /// dropped.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(maybe_uninit_fill)]
+    /// use std::mem::MaybeUninit;
+    ///
+    /// let mut buf = [const { MaybeUninit::<usize>::uninit() }; 5];
+    /// let initialized = buf.write_default();
+    /// assert_eq!(initialized, &mut [0, 0, 0, 0, 0]);
+    /// ```
+    #[unstable(feature = "maybe_uninit_fill", issue = "117428")]
+    pub fn write_default(&mut self) -> &mut [T]
+    where
+        T: Default,
+    {
+        trait DefaultSpec: Default {
+            fn write_default(buf: &mut [MaybeUninit<Self>]) -> &mut [Self];
+        }
+
+        impl<T: Default> DefaultSpec for T {
+            default fn write_default(buf: &mut [MaybeUninit<Self>]) -> &mut [Self] {
+                buf.write_with(|_| T::default())
+            }
+        }
+
+        macro_rules! spec_default_zero {
+            ($ty:ty) => {
+                impl DefaultSpec for $ty {
+                    fn write_default(buf: &mut [MaybeUninit<Self>]) -> &mut [Self] {
+                        // SAFETY:
+                        // `Default::default` is equivalent to zero-initialization
+                        // for all these types, and this initializes the entire
+                        // slice.
+                        unsafe {
+                            buf.as_mut_ptr().write_bytes(0, buf.len());
+                            buf.assume_init_mut()
+                        }
+                    }
+                }
+            };
+        }
+
+        spec_default_zero!(i8);
+        spec_default_zero!(u8);
+        spec_default_zero!(i16);
+        spec_default_zero!(u16);
+        spec_default_zero!(i32);
+        spec_default_zero!(u32);
+        spec_default_zero!(i64);
+        spec_default_zero!(u64);
+        spec_default_zero!(i128);
+        spec_default_zero!(u128);
+        spec_default_zero!(isize);
+        spec_default_zero!(usize);
+
+        T::write_default(self)
+    }
+
     /// Fills a slice with elements yielded by an iterator until either all elements have been
     /// initialized or the iterator is empty.
     ///