1 files changed, 52 insertions, 23 deletions
diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs
index eae0e1c7618..2090756d7a3 100644
--- a/library/core/src/array/mod.rs
+++ b/library/core/src/array/mod.rs
@@ -865,24 +865,6 @@ where
         return Ok(Try::from_output(unsafe { mem::zeroed() }));
     }
 
-    struct Guard<'a, T, const N: usize> {
-        array_mut: &'a mut [MaybeUninit<T>; N],
-        initialized: usize,
-    }
-
-    impl<T, const N: usize> Drop for Guard<'_, T, N> {
-        fn drop(&mut self) {
-            debug_assert!(self.initialized <= N);
-
-            // SAFETY: this slice will contain only initialized objects.
-            unsafe {
-                crate::ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(
-                    &mut self.array_mut.get_unchecked_mut(..self.initialized),
-                ));
-            }
-        }
-    }
-
     let mut array = MaybeUninit::uninit_array::<N>();
     let mut guard = Guard { array_mut: &mut array, initialized: 0 };
 
@@ -896,13 +878,11 @@ where
                     ControlFlow::Continue(elem) => elem,
                 };
 
-                // SAFETY: `guard.initialized` starts at 0, is increased by one in the
-                // loop and the loop is aborted once it reaches N (which is
-                // `array.len()`).
+                // SAFETY: `guard.initialized` starts at 0, which means push can be called
+                // at most N times, which this loop does.
                 unsafe {
-                    guard.array_mut.get_unchecked_mut(guard.initialized).write(item);
+                    guard.push_unchecked(item);
                 }
-                guard.initialized += 1;
             }
             None => {
                 let alive = 0..guard.initialized;
@@ -920,6 +900,55 @@ where
     Ok(Try::from_output(output))
 }
 
+/// Panic guard for incremental initialization of arrays.
+///
+/// Disarm the guard with `mem::forget` once the array has been initialized.
+///
+/// # Safety
+///
+/// All write accesses to this structure are unsafe and must maintain a correct
+/// count of `initialized` elements.
+///
+/// To minimize indirection fields are still pub but callers should at least use
+/// `push_unchecked` to signal that something unsafe is going on.
+pub(crate) struct Guard<'a, T, const N: usize> {
+    /// The array to be initialized.
+    pub array_mut: &'a mut [MaybeUninit<T>; N],
+    /// The number of items that have been initialized so far.
+    pub initialized: usize,
+}
+
+impl<T, const N: usize> Guard<'_, T, N> {
+    /// Adds an item to the array and updates the initialized item counter.
+    ///
+    /// # Safety
+    ///
+    /// No more than N elements must be initialized.
+    #[inline]
+    pub unsafe fn push_unchecked(&mut self, item: T) {
+        // SAFETY: If `initialized` was correct before and the caller does not
+        // invoke this method more than N times then writes will be in-bounds
+        // and slots will not be initialized more than once.
+        unsafe {
+            self.array_mut.get_unchecked_mut(self.initialized).write(item);
+            self.initialized = self.initialized.unchecked_add(1);
+        }
+    }
+}
+
+impl<T, const N: usize> Drop for Guard<'_, T, N> {
+    fn drop(&mut self) {
+        debug_assert!(self.initialized <= N);
+
+        // SAFETY: this slice will contain only initialized objects.
+        unsafe {
+            crate::ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(
+                &mut self.array_mut.get_unchecked_mut(..self.initialized),
+            ));
+        }
+    }
+}
+
 /// Returns the next chunk of `N` items from the iterator or errors with an
 /// iterator over the remainder. Used for `Iterator::next_chunk`.
 #[inline]