diff --git a/library/core/src/iter/adapters/map_windows.rs b/library/core/src/iter/adapters/map_windows.rs index 097a0745c61c7..727784bdd50c6 100644 --- a/library/core/src/iter/adapters/map_windows.rs +++ b/library/core/src/iter/adapters/map_windows.rs @@ -1,5 +1,5 @@ use crate::iter::FusedIterator; -use crate::mem::MaybeUninit; +use crate::mem::{ManuallyDrop, MaybeUninit}; use crate::{fmt, ptr}; /// An iterator over the mapped windows of another iterator. @@ -31,15 +31,27 @@ struct MapWindowsInner { } // `Buffer` uses two times of space to reduce moves among the iterations. +struct Buffer { + // Invariant: N elements starting from `self.start` must be initialized at + // with all other elements left uninitialized. + buffer: RawBuffer, + // Invariant: `self.start <= N` + start: usize, +} + +/// Internal storage for `Buffer`. +/// +/// Has no storage invariants, but has access invariants. +/// See the unsafe method contracts for more information. +/// +/// This type does not implement Drop, it will leak any data stored within. +// // `Buffer` is semantically `[MaybeUninit; 2 * N]`. However, due // to limitations of const generics, we use this different type. Note that // it has the same underlying memory layout. -struct Buffer { - // Invariant: `self.buffer[self.start..self.start + N]` is initialized, - // with all other elements being uninitialized. This also - // implies that `self.start <= N`. - buffer: [[MaybeUninit; N]; 2], - start: usize, +#[repr(transparent)] +struct RawBuffer { + data: [[MaybeUninit; N]; 2], } impl MapWindows { @@ -79,7 +91,16 @@ impl MapWindowsInner { Some(item) => buffer.push(item), }, } - self.buffer.as_ref().map(Buffer::as_array_ref) + self.buffer.as_ref().map(|buf| + // SAFETY: + // - if this was the first time to advance, this is a new well-formed `Buffer`. + // + // - if we already had a buffer, and `iter.next()` was Some; + // `Buffer::push` is responsible for upholding the invariant before we reach this. + // + // - if we already had a buffer, and `iter.next()` was None; + // this closure is unreachable, as the buffer was taken beforehand. + unsafe { buf.buffer.as_array_ref(buf.start) }) } fn size_hint(&self) -> (usize, Option) { @@ -98,37 +119,24 @@ impl MapWindowsInner { } impl Buffer { - fn try_from_iter(iter: &mut impl Iterator) -> Option { - let first_half = crate::array::iter_next_chunk(iter).ok()?; - let buffer = - [MaybeUninit::new(first_half).transpose(), [const { MaybeUninit::uninit() }; N]]; - Some(Self { buffer, start: 0 }) - } - - #[inline] - fn buffer_ptr(&self) -> *const MaybeUninit { - self.buffer.as_ptr().cast() - } - - #[inline] - fn buffer_mut_ptr(&mut self) -> *mut MaybeUninit { - self.buffer.as_mut_ptr().cast() - } - + /// # Safety + /// + /// This type implements `Drop`, and it has invariants that must be upheld: + /// - `start` must be within the bounds `0..=N`. + /// - `raw[start..start + N]` must be initialized. #[inline] - fn as_array_ref(&self) -> &[T; N] { - debug_assert!(self.start + N <= 2 * N); - - // SAFETY: our invariant guarantees these elements are initialized. - unsafe { &*self.buffer_ptr().add(self.start).cast() } + unsafe fn new(raw: RawBuffer, start: usize) -> Self { + Self { buffer: raw, start } } - #[inline] - fn as_uninit_array_mut(&mut self) -> &mut MaybeUninit<[T; N]> { - debug_assert!(self.start + N <= 2 * N); - - // SAFETY: our invariant guarantees these elements are in bounds. - unsafe { &mut *self.buffer_mut_ptr().add(self.start).cast() } + fn try_from_iter(iter: &mut impl Iterator) -> Option { + let first_half: [T; N] = crate::array::iter_next_chunk(iter).ok()?; + let raw = RawBuffer { + data: [MaybeUninit::new(first_half).transpose(), [const { MaybeUninit::uninit() }; N]], + }; + // SAFETY: buffer has the first `first_half.len()` items initialized, which upholds the + // internal invariant for the start offset 0, which is also within bounds. + Some(unsafe { Self::new(raw, 0) }) } /// Pushes a new item `next` to the back, and pops the front-most one. @@ -136,7 +144,7 @@ impl Buffer { /// All the elements will be shifted to the front end when pushing reaches /// the back end. fn push(&mut self, next: T) { - let buffer_mut_ptr = self.buffer_mut_ptr(); + let buffer_mut_ptr = self.buffer.as_mut_ptr(); debug_assert!(self.start + N <= 2 * N); let to_drop = if self.start == N { @@ -188,18 +196,77 @@ impl Buffer { // SAFETY: the index is valid and this is element `a` in the // diagram above and has not been dropped yet. - unsafe { ptr::drop_in_place(to_drop.cast_init()) }; + unsafe { to_drop.cast_init().drop_in_place() }; + } +} + +impl RawBuffer { + #[inline] + fn as_ptr(&self) -> *const MaybeUninit { + self.data.as_ptr().cast() + } + + #[inline] + fn as_mut_ptr(&mut self) -> *mut MaybeUninit { + self.data.as_mut_ptr().cast() + } + + /// # Safety + /// + /// `self.data` must uphold the internal invariants of `Buffer`: + /// - `start` must be within the bounds `0..=N` + /// - `self.data[start..start + N]` must be initialized. + #[inline] + unsafe fn as_array_ref(&self, start: usize) -> &[T; N] { + debug_assert!(start + N <= 2 * N); + + // SAFETY: caller satisfies that `start` is within bounds. + unsafe { &*self.as_ptr().add(start).cast() } + } + + /// # Safety + /// + /// `self.data` must uphold the internal invariants of `Buffer`: + /// - `start` must be within the bounds `0..=N` + /// - `self.data[start..start + N]` must not overlap with the initialized part. + #[inline] + unsafe fn as_uninit_array_mut(&mut self, start: usize) -> &mut MaybeUninit<[T; N]> { + debug_assert!(start + N <= 2 * N); + + // SAFETY: caller satisfies that `start` is within bounds. + unsafe { &mut *self.as_mut_ptr().add(start).cast() } } } impl Clone for Buffer { fn clone(&self) -> Self { - let mut buffer = Buffer { - buffer: [[const { MaybeUninit::uninit() }; N], [const { MaybeUninit::uninit() }; N]], + let mut md = ManuallyDrop::new(Buffer { + // INVARIANT: the `Buffer` is transiently wrapped in ManuallyDrop so + // the it is never actually dropped. + buffer: RawBuffer { + data: [ + [const { MaybeUninit::::uninit() }; N], + [const { MaybeUninit::::uninit() }; N], + ], + }, + // INVARIANT: our invariants guarantee `self.start` was within bounds, + // so this must be too. start: self.start, - }; - buffer.as_uninit_array_mut().write(self.as_array_ref().clone()); - buffer + }); + + // SAFETY: `md.buffer` is currently fully uninitialized, which can not + // overlap with any initialized part. + // Further, `self`'s invariants guarantee N elements starting at + // `self.start` are initialized within `self.buffer`. + unsafe { + md.buffer.as_uninit_array_mut(self.start).write( + // if clone starts unwinding, the `Buffer` is not dropped due + // to being inside a ManuallyDrop. + self.buffer.as_array_ref(self.start).clone(), + ); + } + + ManuallyDrop::into_inner(md) } } @@ -215,15 +282,12 @@ where impl Drop for Buffer { fn drop(&mut self) { - // SAFETY: our invariant guarantees that N elements starting from - // `self.start` are initialized. We drop them here. - unsafe { - let initialized_part: *mut [T] = crate::ptr::slice_from_raw_parts_mut( - self.buffer_mut_ptr().add(self.start).cast(), - N, - ); - ptr::drop_in_place(initialized_part); - } + // SAFETY: our invariant guarantees that `self.start` is within bounds + let init_ptr = unsafe { self.buffer.as_mut_ptr().add(self.start).cast_init() }; + + // SAFETY: our invariant guarantees N elements starting from + // `init_ptr` are initialized. We drop them here. + unsafe { init_ptr.cast_slice(N).drop_in_place() }; } }