From 2419cb5faa7ef0bd07b8a1e898aa4eaa20452f83 Mon Sep 17 00:00:00 2001
From: Roman Proskuryakov <r.proskuryakoff@gmail.com>
Date: Sat, 11 Sep 2021 00:41:41 +0300
Subject: [PATCH] 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

---
 library/core/src/mem/maybe_uninit.rs | 54 ++++++++++++++++++++++++++++
 library/core/tests/lib.rs            |  1 +
 library/core/tests/mem.rs            | 50 ++++++++++++++++++++++++++
 3 files changed, 105 insertions(+)

diff --git a/library/core/src/mem/maybe_uninit.rs b/library/core/src/mem/maybe_uninit.rs
index 9c88a62336161..102d9d157cc0d 100644
--- a/library/core/src/mem/maybe_uninit.rs
+++ b/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()
+    }
+}
diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs
index 19bcc45108dfd..fd790036becba 100644
--- a/library/core/tests/lib.rs
+++ b/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)]
diff --git a/library/core/tests/mem.rs b/library/core/tests/mem.rs
index c780bb32ca96d..5b148306aa04c 100644
--- a/library/core/tests/mem.rs
+++ b/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]);
+}