From 3925fc0c8e0b4ad9be48fbc1136281841d8985fa Mon Sep 17 00:00:00 2001
From: The 8472 <git@infinite-source.de>
Date: Tue, 8 Nov 2022 00:13:26 +0100
Subject: [PATCH] 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.
---
 library/core/src/array/mod.rs                 | 39 ++++++++++++++++---
 .../core/src/iter/adapters/array_chunks.rs    |  8 ++--
 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