diff options
| -rw-r--r-- | library/core/src/array/mod.rs | 39 | ||||
| -rw-r--r-- | library/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 |