Add impl Index for MaybeUninit<[T; N]>

Add impl IndexMut for `MaybeUninit<[T; N]>`
Add impl IntoIterator for `&mut MaybeUninit<[T; N]>`
Add `MaybeUninit<[T; N]>::iter_mut` method

Author: kpp

library/core/src/mem/maybe_uninit.rs

@@ -2,7 +2,9 @@ use crate::any::type_name;
 use crate::fmt;
 use crate::intrinsics;
 use crate::mem::ManuallyDrop;
+use crate::ops::{Index, IndexMut};
 use crate::ptr;
+use crate::slice::IterMut;
 
 /// A wrapper type to construct uninitialized instances of `T`.
 ///
@@ -1158,3 +1160,55 @@ impl<T> MaybeUninit<T> {
         unsafe { MaybeUninit::slice_assume_init_mut(this) }
     }
 }
+
+impl<T, const N: usize> Index<usize> for MaybeUninit<[T; N]> {
+    type Output = MaybeUninit<T>;
+
+    #[inline]
+    fn index(&self, index: usize) -> &Self::Output {
+        Index::index(self.as_ref(), index)
+    }
+}
+
+impl<T, const N: usize> IndexMut<usize> for MaybeUninit<[T; N]> {
+    #[inline]
+    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
+        IndexMut::index_mut(self.as_mut(), index)
+    }
+}
+
+impl<T, const N: usize> AsRef<[MaybeUninit<T>]> for MaybeUninit<[T; N]> {
+    #[inline]
+    fn as_ref(&self) -> &[MaybeUninit<T>] {
+        let data = self.as_ptr().cast::<MaybeUninit<T>>();
+        // SAFETY: MaybeUninit<[T; N]> and [MaybeUninit<T>; N] have the same layout
+        // data points to N consecutive properly initialized values of type MaybeUninit<T>.
+        unsafe { crate::slice::from_raw_parts(data, N) }
+    }
+}
+
+impl<T, const N: usize> AsMut<[MaybeUninit<T>]> for MaybeUninit<[T; N]> {
+    #[inline]
+    fn as_mut(&mut self) -> &mut [MaybeUninit<T>] {
+        let data = self.as_mut_ptr().cast::<MaybeUninit<T>>();
+        // SAFETY: MaybeUninit<[T; N]> and [MaybeUninit<T>; N] have the same layout
+        // data points to N consecutive properly initialized values of type MaybeUninit<T>.
+        unsafe { crate::slice::from_raw_parts_mut(data, N) }
+    }
+}
+
+impl<'a, T, const N: usize> IntoIterator for &'a mut MaybeUninit<[T; N]> {
+    type Item = &'a mut MaybeUninit<T>;
+    type IntoIter = IterMut<'a, MaybeUninit<T>>;
+
+    fn into_iter(self) -> IterMut<'a, MaybeUninit<T>> {
+        self.as_mut().iter_mut()
+    }
+}
+
+#[unstable(feature = "maybe_uninit_array_index", issue = "none")]
+impl<T, const N: usize> MaybeUninit<[T; N]> {
+    pub fn iter_mut(&mut self) -> IterMut<'_, MaybeUninit<T>> {
+        self.as_mut().iter_mut()
+    }
+}

library/core/tests/lib.rs

@@ -34,6 +34,7 @@
 #![feature(slice_partition_at_index)]
 #![feature(maybe_uninit_uninit_array)]
 #![feature(maybe_uninit_array_assume_init)]
+#![feature(maybe_uninit_array_index)]
 #![feature(maybe_uninit_extra)]
 #![feature(maybe_uninit_write_slice)]
 #![feature(min_specialization)]

library/core/tests/mem.rs

@@ -269,3 +269,53 @@ fn uninit_const_assume_init_read() {
     const FOO: u32 = unsafe { MaybeUninit::new(42).assume_init_read() };
     assert_eq!(FOO, 42);
 }
+
+#[test]
+fn uninit_array_index_mut() {
+    let mut v = MaybeUninit::<[i32; 4]>::uninit();
+
+    v[0].write(1);
+    v[1].write(2);
+    v[2].write(4);
+    v[3].write(8);
+
+    let v = unsafe { v.assume_init() };
+    assert_eq!(v, [1, 2, 4, 8]);
+}
+
+#[test]
+fn uninit_array_as_mut() {
+    let mut v = MaybeUninit::<[i32; 4]>::uninit();
+
+    v.as_mut()[0].write(1);
+    v.as_mut()[1].write(2);
+    v.as_mut()[2].write(4);
+    v.as_mut()[3].write(8);
+
+    let v = unsafe { v.assume_init() };
+    assert_eq!(v, [1, 2, 4, 8]);
+}
+
+#[test]
+fn uninit_array_loop_init() {
+    let mut v = MaybeUninit::<[i32; 4]>::uninit();
+
+    for v in &mut v {
+        v.write(42);
+    }
+
+    let v = unsafe { v.assume_init() };
+    assert_eq!(v, [42, 42, 42, 42]);
+}
+
+#[test]
+fn uninit_array_iter_mut() {
+    let mut v = MaybeUninit::<[i32; 4]>::uninit();
+
+    v.iter_mut().zip(1..=4).for_each(|(v, val)| {
+        v.write(val);
+    });
+
+    let v = unsafe { v.assume_init() };
+    assert_eq!(v, [1, 2, 4, 8]);
+}