Skip to content

Commit 3925fc0

Browse files
committed
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.
1 parent 43c353f commit 3925fc0

File tree

2 files changed

+39
-8
lines changed

2 files changed

+39
-8
lines changed

library/core/src/array/mod.rs

Lines changed: 34 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -878,13 +878,11 @@ where
878878
ControlFlow::Continue(elem) => elem,
879879
};
880880

881-
// SAFETY: `guard.initialized` starts at 0, is increased by one in the
882-
// loop and the loop is aborted once it reaches N (which is
883-
// `array.len()`).
881+
// SAFETY: `guard.initialized` starts at 0, which means push can be called
882+
// at most N times, which this loop does.
884883
unsafe {
885-
guard.array_mut.get_unchecked_mut(guard.initialized).write(item);
884+
guard.push_unchecked(item);
886885
}
887-
guard.initialized += 1;
888886
}
889887
None => {
890888
let alive = 0..guard.initialized;
@@ -902,11 +900,42 @@ where
902900
Ok(Try::from_output(output))
903901
}
904902

903+
/// Panic guard for incremental initialization of arrays.
904+
///
905+
/// Disarm the guard with `mem::forget` once the array has been initialized.
906+
///
907+
/// # Safety
908+
///
909+
/// All write accesses to this structure are unsafe and must maintain a correct
910+
/// count of `initialized` elements.
911+
///
912+
/// To minimize indirection fields are still pub but callers should at least use
913+
/// `push_unchecked` to signal that something unsafe is going on.
905914
pub(crate) struct Guard<'a, T, const N: usize> {
915+
/// The array to be initialized.
906916
pub array_mut: &'a mut [MaybeUninit<T>; N],
917+
/// The number of items that have been initialized so far.
907918
pub initialized: usize,
908919
}
909920

921+
impl<T, const N: usize> Guard<'_, T, N> {
922+
/// Adds an item to the array and updates the initialized item counter.
923+
///
924+
/// # Safety
925+
///
926+
/// No more than N elements must be initialized.
927+
#[inline]
928+
pub unsafe fn push_unchecked(&mut self, item: T) {
929+
// SAFETY: If `initialized` was correct before and the caller does not
930+
// invoke this method more than N times then writes will be in-bounds
931+
// and slots will not be initialized more than once.
932+
unsafe {
933+
self.array_mut.get_unchecked_mut(self.initialized).write(item);
934+
self.initialized = self.initialized.unchecked_add(1);
935+
}
936+
}
937+
}
938+
910939
impl<T, const N: usize> Drop for Guard<'_, T, N> {
911940
fn drop(&mut self) {
912941
debug_assert!(self.initialized <= N);

library/core/src/iter/adapters/array_chunks.rs

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -216,11 +216,13 @@ where
216216
while inner_len - i >= N {
217217
let mut chunk = MaybeUninit::uninit_array();
218218
let mut guard = array::Guard { array_mut: &mut chunk, initialized: 0 };
219-
for j in 0..N {
219+
while guard.initialized < N {
220220
// SAFETY: The method consumes the iterator and the loop condition ensures that
221221
// all accesses are in bounds and only happen once.
222-
guard.array_mut[j].write(unsafe { self.iter.__iterator_get_unchecked(i + j) });
223-
guard.initialized = j + 1;
222+
unsafe {
223+
let idx = i + guard.initialized;
224+
guard.push_unchecked(self.iter.__iterator_get_unchecked(idx));
225+
}
224226
}
225227
mem::forget(guard);
226228
// SAFETY: The loop above initialized all elements

0 commit comments

Comments
 (0)