What if we just add <[MaybeUninit<T>; N]>::assume_init directly?

Author: scottmcm

library/alloc/src/lib.rs

@@ -92,6 +92,7 @@
 #![feature(allocator_api)]
 #![feature(array_chunks)]
 #![feature(array_into_iter_constructors)]
+#![feature(array_maybe_uninit_assume_init)]
 #![feature(array_methods)]
 #![feature(array_windows)]
 #![feature(assert_matches)]
@@ -127,7 +128,6 @@
 #![feature(layout_for_ptr)]
 #![feature(maybe_uninit_slice)]
 #![feature(maybe_uninit_uninit_array)]
-#![feature(maybe_uninit_uninit_array_transpose)]
 #![cfg_attr(test, feature(new_uninit))]
 #![feature(nonnull_slice_from_raw_parts)]
 #![feature(pattern)]

library/alloc/src/vec/into_iter.rs

@@ -223,7 +223,7 @@ impl<T, A: Allocator> Iterator for IntoIter<T, A> {
 
             self.ptr = self.ptr.wrapping_byte_add(N);
             // Safety: ditto
-            return Ok(unsafe { raw_ary.transpose().assume_init() });
+            return Ok(unsafe { raw_ary.assume_init() });
         }
 
         if len < N {
@@ -241,7 +241,7 @@ impl<T, A: Allocator> Iterator for IntoIter<T, A> {
         return unsafe {
             ptr::copy_nonoverlapping(self.ptr, raw_ary.as_mut_ptr() as *mut T, N);
             self.ptr = self.ptr.add(N);
-            Ok(raw_ary.transpose().assume_init())
+            Ok(raw_ary.assume_init())
         };
     }
 

library/core/src/array/iter.rs

@@ -104,7 +104,7 @@ impl<T, const N: usize> IntoIter<T, N> {
     ///
     /// ```
     /// #![feature(array_into_iter_constructors)]
-    /// #![feature(maybe_uninit_uninit_array_transpose)]
+    /// #![feature(array_maybe_uninit_assume_init)]
     /// #![feature(maybe_uninit_uninit_array)]
     /// use std::array::IntoIter;
     /// use std::mem::MaybeUninit;
@@ -133,7 +133,7 @@ impl<T, const N: usize> IntoIter<T, N> {
     ///     }
     ///
     ///     // SAFETY: We've initialized all N items
-    ///     unsafe { Ok(buffer.transpose().assume_init()) }
+    ///     unsafe { Ok(buffer.assume_init()) }
     /// }
     ///
     /// let r: [_; 4] = next_chunk(&mut (10..16)).unwrap();

library/core/src/array/mod.rs

@@ -799,6 +799,51 @@ impl<T, const N: usize> [T; N] {
     }
 }
 
+impl<T, const N: usize> [MaybeUninit<T>; N] {
+    /// Extracts the values from an array of [`MaybeUninit<T>`] containers.
+    ///
+    /// This is essentially `.map(MaybeUninit::assume_init)`, but in a way
+    /// that's *much* clearer to the optimizer how to do efficiently.
+    ///
+    /// # Safety
+    ///
+    /// It is up to the caller to guarantee that all elements of the array are
+    /// in an initialized state.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(inline_const)]
+    /// #![feature(array_maybe_uninit_assume_init)]
+    /// use std::mem::MaybeUninit;
+    ///
+    /// let mut array = [const { MaybeUninit::<i32>::uninit() }; 3];
+    /// array[0].write(0);
+    /// array[1].write(1);
+    /// array[2].write(2);
+    ///
+    /// // SAFETY: we initialised all three elements
+    /// let array = unsafe { array.assume_init() };
+    ///
+    /// assert_eq!(array, [0, 1, 2]);
+    /// ```
+    #[unstable(feature = "array_maybe_uninit_assume_init", issue = "96097")]
+    #[rustc_const_unstable(feature = "const_array_maybe_uninit_assume_init", issue = "96097")]
+    #[inline(always)]
+    #[track_caller]
+    pub const unsafe fn assume_init(self) -> [T; N] {
+        // SAFETY:
+        // * The caller guarantees that all elements of the array are initialized
+        // * `MaybeUninit<T>` and T are guaranteed to have the same layout
+        // * `MaybeUninit` does not drop, so there are no double-frees
+        // And thus the conversion is safe
+        unsafe {
+            crate::intrinsics::assert_inhabited::<[T; N]>();
+            mem::transmute_copy(&mem::ManuallyDrop::new(self))
+        }
+    }
+}
+
 /// Pulls `N` items from `iter` and returns them as an array. If the iterator
 /// yields fewer than `N` items, this function exhibits undefined behavior.
 ///
@@ -896,7 +941,7 @@ where
 
     mem::forget(guard);
     // SAFETY: All elements of the array were populated in the loop above.
-    let output = unsafe { array.transpose().assume_init() };
+    let output = unsafe { array.assume_init() };
     Ok(Try::from_output(output))
 }
 

library/core/tests/lib.rs

@@ -1,5 +1,6 @@
 #![feature(alloc_layout_extra)]
 #![feature(array_chunks)]
+#![feature(array_maybe_uninit_assume_init)]
 #![feature(array_methods)]
 #![feature(array_windows)]
 #![feature(bigint_helper_methods)]
@@ -53,7 +54,6 @@
 #![feature(split_as_slice)]
 #![feature(maybe_uninit_uninit_array)]
 #![feature(maybe_uninit_write_slice)]
-#![feature(maybe_uninit_uninit_array_transpose)]
 #![feature(min_specialization)]
 #![feature(numfmt)]
 #![feature(step_trait)]

library/core/tests/mem.rs

@@ -194,11 +194,11 @@ fn uninit_array_assume_init() {
     array[3].write(1);
     array[4].write(5);
 
-    let array = unsafe { array.transpose().assume_init() };
+    let array = unsafe { array.assume_init() };
 
     assert_eq!(array, [3, 1, 4, 1, 5]);
 
-    let [] = unsafe { [MaybeUninit::<!>::uninit(); 0].transpose().assume_init() };
+    let [] = unsafe { [MaybeUninit::<!>::uninit(); 0].assume_init() };
 }
 
 #[test]