Refactor underlying data of map_windows::Buffer into RawBuffer

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<I: Iterator, const N: usize> {
 }
 
 // `Buffer` uses two times of space to reduce moves among the iterations.
+struct Buffer<T, const N: usize> {
+    // Invariant: N elements starting from `self.start` must be initialized at
+    // with all other elements left uninitialized.
+    buffer: RawBuffer<T, N>,
+    // Invariant: `self.start <= N`
+    start: usize,
+}
+
+/// Internal storage for `Buffer<T, N>`.
+///
+/// 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<T, N>` is semantically `[MaybeUninit<T>; 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<T, const N: usize> {
-    // 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<T>; N]; 2],
-    start: usize,
+#[repr(transparent)]
+struct RawBuffer<T, const N: usize> {
+    data: [[MaybeUninit<T>; N]; 2],
 }
 
 impl<I: Iterator, F, const N: usize> MapWindows<I, F, N> {
@@ -79,7 +91,16 @@ impl<I: Iterator, const N: usize> MapWindowsInner<I, N> {
                 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<usize>) {
@@ -98,45 +119,32 @@ impl<I: Iterator, const N: usize> MapWindowsInner<I, N> {
 }
 
 impl<T, const N: usize> Buffer<T, N> {
-    fn try_from_iter(iter: &mut impl Iterator<Item = T>) -> Option<Self> {
-        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<T> {
-        self.buffer.as_ptr().cast()
-    }
-
-    #[inline]
-    fn buffer_mut_ptr(&mut self) -> *mut MaybeUninit<T> {
-        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<T, N>, 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<Item = T>) -> Option<Self> {
+        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.
     ///
     /// 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<T, const N: usize> Buffer<T, N> {
 
         // 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<T, const N: usize> RawBuffer<T, N> {
+    #[inline]
+    fn as_ptr(&self) -> *const MaybeUninit<T> {
+        self.data.as_ptr().cast()
+    }
+
+    #[inline]
+    fn as_mut_ptr(&mut self) -> *mut MaybeUninit<T> {
+        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<T: Clone, const N: usize> Clone for Buffer<T, N> {
     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::<T>::uninit() }; N],
+                    [const { MaybeUninit::<T>::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<T, const N: usize> Drop for Buffer<T, N> {
     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() };
     }
 }