document and improve array Guard type

The type is unsafe and now exposed to the whole crate.
Document it properly and add an unsafe method so the
caller can make it visible that something unsafe is happening.

2 files changed, 39 insertions, 8 deletions

diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs
index 24e5d0fba21..2090756d7a3 100644
--- a/library/core/src/array/mod.rs
+++ b/library/core/src/array/mod.rs
@@ -878,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;
@@ -902,11 +900,42 @@ 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);
diff --git a/library/core/src/iter/adapters/array_chunks.rs b/library/core/src/iter/adapters/array_chunks.rs
index 9bd135007b4..5e4211058aa 100644
--- a/library/core/src/iter/adapters/array_chunks.rs
+++ b/library/core/src/iter/adapters/array_chunks.rs
@@ -216,11 +216,13 @@ where
         while inner_len - i >= N {
             let mut chunk = MaybeUninit::uninit_array();
             let mut guard = array::Guard { array_mut: &mut chunk, initialized: 0 };
-            for j in 0..N {
+            while guard.initialized < N {
                 // SAFETY: The method consumes the iterator and the loop condition ensures that
                 // all accesses are in bounds and only happen once.
-                guard.array_mut[j].write(unsafe { self.iter.__iterator_get_unchecked(i + j) });
-                guard.initialized = j + 1;
+                unsafe {
+                    let idx = i + guard.initialized;
+                    guard.push_unchecked(self.iter.__iterator_get_unchecked(idx));
+                }
             }
             mem::forget(guard);
             // SAFETY: The loop above initialized all elements