map_windows.rs - source
Description: Source of the Rust file `library/core/src/iter/adapters/map_windows.rs`.

#
core/iter/adapters/
map_windows.rs

    1use crate::iter::FusedIterator;
    2use crate::mem::MaybeUninit;
    3use crate::{fmt, ptr};
    4
    5/// An iterator over the mapped windows of another iterator.
    6///
    7/// This `struct` is created by the [`Iterator::map_windows`]. See its
    8/// documentation for more information.
    9#[must_use = "iterators are lazy and do nothing unless consumed"]
    10#[unstable(feature = "iter_map_windows", issue = "87155")]
    11pub struct MapWindows<I: Iterator, F, const N: usize> {
    12    f: F,
    13    inner: MapWindowsInner<I, N>,
    14}
    15
    16struct MapWindowsInner<I: Iterator, const N: usize> {
    17    // We fuse the inner iterator because there shouldn't be "holes" in
    18    // the sliding window. Once the iterator returns a `None`, we make
    19    // our `MapWindows` iterator return `None` forever.
    20    iter: Option<I>,
    21    // Since iterators are assumed lazy, i.e. it only yields an item when
    22    // `Iterator::next()` is called, and `MapWindows` is not an exception.
    23    //
    24    // Before the first iteration, we keep the buffer `None`. When the user
    25    // first call `next` or other methods that makes the iterator advance,
    26    // we collect the first `N` items yielded from the inner iterator and
    27    // put it into the buffer.
    28    //
    29    // When the inner iterator has returned a `None` (i.e. fused), we take
    30    // away this `buffer` and leave it `None` to reclaim its resources.
    31    //
    32    // FIXME: should we shrink the size of `buffer` using niche optimization?
    33    buffer: Option<Buffer<I::Item, N>>,
    34}
    35
    36// `Buffer` uses two times of space to reduce moves among the iterations.
    37// `Buffer<T, N>` is semantically `[MaybeUninit<T>; 2 * N]`. However, due
    38// to limitations of const generics, we use this different type. Note that
    39// it has the same underlying memory layout.
    40struct Buffer<T, const N: usize> {
    41    // Invariant: `self.buffer[self.start..self.start + N]` is initialized,
    42    // with all other elements being uninitialized. This also
    43    // implies that `self.start <= N`.
    44    buffer: [[MaybeUninit<T>; N]; 2],
    45    start: usize,
    46}
    47
    48impl<I: Iterator, F, const N: usize> MapWindows<I, F, N> {
    49    pub(in crate::iter) fn new(iter: I, f: F) -> Self {
    50        assert!(N != 0, "array in `Iterator::map_windows` must contain more than 0 elements");
    51
    52        // Only ZST arrays' length can be so large.
    53        if size_of::<I::Item>() == 0 {
    54            assert!(
    55                N.checked_mul(2).is_some(),
    56                "array size of `Iterator::map_windows` is too large"
    57            );
    58        }
    59
    60        Self { inner: MapWindowsInner::new(iter), f }
    61    }
    62}
    63
    64impl<I: Iterator, const N: usize> MapWindowsInner<I, N> {
    65    #[inline]
    66    fn new(iter: I) -> Self {
    67        Self { iter: Some(iter), buffer: None }
    68    }
    69
    70    fn next_window(&mut self) -> Option<&[I::Item; N]> {
    71        let iter = self.iter.as_mut()?;
    72        match self.buffer {
    73            // It is the first time to advance. We collect
    74            // the first `N` items from `self.iter` to initialize `self.buffer`.
    75            None => self.buffer = Buffer::try_from_iter(iter),
    76            Some(ref mut buffer) => match iter.next() {
    77                None => {
    78                    // Fuse the inner iterator since it yields a `None`.
    79                    self.iter.take();
    80                    self.buffer.take();
    81                }
    82                // Advance the iterator. We first call `next` before changing our buffer
    83                // at all. This means that if `next` panics, our invariant is upheld and
    84                // our `Drop` impl drops the correct elements.
    85                Some(item) => buffer.push(item),
    86            },
    87        }
    88        self.buffer.as_ref().map(Buffer::as_array_ref)
    89    }
    90
    91    fn size_hint(&self) -> (usize, Option<usize>) {
    92        let Some(ref iter) = self.iter else { return (0, Some(0)) };
    93        let (lo, hi) = iter.size_hint();
    94        if self.buffer.is_some() {
    95            // If the first `N` items are already yielded by the inner iterator,
    96            // the size hint is then equal to the that of the inner iterator's.
    97            (lo, hi)
    98        } else {
    99            // If the first `N` items are not yet yielded by the inner iterator,
    100            // the first `N` elements should be counted as one window, so both bounds
    101            // should subtract `N - 1`.
    102            (lo.saturating_sub(N - 1), hi.map(|hi| hi.saturating_sub(N - 1)))
    103        }
    104    }
    105}
    106
    107impl<T, const N: usize> Buffer<T, N> {
    108    fn try_from_iter(iter: &mut impl Iterator<Item = T>) -> Option<Self> {
    109        let first_half = crate::array::iter_next_chunk(iter).ok()?;
    110        let buffer =
    111            [MaybeUninit::new(first_half).transpose(), [const { MaybeUninit::uninit() }; N]];
    112        Some(Self { buffer, start: 0 })
    113    }
    114
    115    #[inline]
    116    fn buffer_ptr(&self) -> *const MaybeUninit<T> {
    117        self.buffer.as_ptr().cast()
    118    }
    119
    120    #[inline]
    121    fn buffer_mut_ptr(&mut self) -> *mut MaybeUninit<T> {
    122        self.buffer.as_mut_ptr().cast()
    123    }
    124
    125    #[inline]
    126    fn as_array_ref(&self) -> &[T; N] {
    127        debug_assert!(self.start + N <= 2 * N);
    128
    129        // SAFETY: our invariant guarantees these elements are initialized.
    130        unsafe { &*self.buffer_ptr().add(self.start).cast() }
    131    }
    132
    133    #[inline]
    134    fn as_uninit_array_mut(&mut self) -> &mut MaybeUninit<[T; N]> {
    135        debug_assert!(self.start + N <= 2 * N);
    136
    137        // SAFETY: our invariant guarantees these elements are in bounds.
    138        unsafe { &mut *self.buffer_mut_ptr().add(self.start).cast() }
    139    }
    140
    141    /// Pushes a new item `next` to the back, and pops the front-most one.
    142    ///
    143    /// All the elements will be shifted to the front end when pushing reaches
    144    /// the back end.
    145    fn push(&mut self, next: T) {
    146        let buffer_mut_ptr = self.buffer_mut_ptr();
    147        debug_assert!(self.start + N <= 2 * N);
    148
    149        let to_drop = if self.start == N {
    150            // We have reached the end of our buffer and have to copy
    151            // everything to the start. Example layout for N = 3.
    152            //
    153            //    0   1   2   3   4   5            0   1   2   3   4   5
    154            //  ┌───┬───┬───┬───┬───┬───┐        ┌───┬───┬───┬───┬───┬───┐
    155            //  │ - │ - │ - │ a │ b │ c │   ->   │ b │ c │ n │ - │ - │ - │
    156            //  └───┴───┴───┴───┴───┴───┘        └───┴───┴───┴───┴───┴───┘
    157            //                ↑                    ↑
    158            //              start                start
    159
    160            // SAFETY: the two pointers are valid for reads/writes of N -1
    161            // elements because our array's size is semantically 2 * N. The
    162            // regions also don't overlap for the same reason.
    163            //
    164            // We leave the old elements in place. As soon as `start` is set
    165            // to 0, we treat them as uninitialized and treat their copies
    166            // as initialized.
    167            let to_drop = unsafe {
    168                ptr::copy_nonoverlapping(buffer_mut_ptr.add(self.start + 1), buffer_mut_ptr, N - 1);
    169                (*buffer_mut_ptr.add(N - 1)).write(next);
    170                buffer_mut_ptr.add(self.start)
    171            };
    172            self.start = 0;
    173            to_drop
    174        } else {
    175            // SAFETY: `self.start` is < N as guaranteed by the invariant
    176            // plus the check above. Even if the drop at the end panics,
    177            // the invariant is upheld.
    178            //
    179            // Example layout for N = 3:
    180            //
    181            //    0   1   2   3   4   5            0   1   2   3   4   5
    182            //  ┌───┬───┬───┬───┬───┬───┐        ┌───┬───┬───┬───┬───┬───┐
    183            //  │ - │ a │ b │ c │ - │ - │   ->   │ - │ - │ b │ c │ n │ - │
    184            //  └───┴───┴───┴───┴───┴───┘        └───┴───┴───┴───┴───┴───┘
    185            //        ↑                                    ↑
    186            //      start                                start
    187            //
    188            let to_drop = unsafe {
    189                (*buffer_mut_ptr.add(self.start + N)).write(next);
    190                buffer_mut_ptr.add(self.start)
    191            };
    192            self.start += 1;
    193            to_drop
    194        };
    195
    196        // SAFETY: the index is valid and this is element `a` in the
    197        // diagram above and has not been dropped yet.
    198        unsafe { ptr::drop_in_place(to_drop.cast_init()) };
    199    }
    200}
    201
    202impl<T: Clone, const N: usize> Clone for Buffer<T, N> {
    203    fn clone(&self) -> Self {
    204        let mut buffer = Buffer {
    205            buffer: [[const { MaybeUninit::uninit() }; N], [const { MaybeUninit::uninit() }; N]],
    206            start: self.start,
    207        };
    208        buffer.as_uninit_array_mut().write(self.as_array_ref().clone());
    209        buffer
    210    }
    211}
    212
    213impl<I, const N: usize> Clone for MapWindowsInner<I, N>
    214where
    215    I: Iterator + Clone,
    216    I::Item: Clone,
    217{
    218    fn clone(&self) -> Self {
    219        Self { iter: self.iter.clone(), buffer: self.buffer.clone() }
    220    }
    221}
    222
    223impl<T, const N: usize> Drop for Buffer<T, N> {
    224    fn drop(&mut self) {
    225        // SAFETY: our invariant guarantees that N elements starting from
    226        // `self.start` are initialized. We drop them here.
    227        unsafe {
    228            let initialized_part: *mut [T] = crate::ptr::slice_from_raw_parts_mut(
    229                self.buffer_mut_ptr().add(self.start).cast(),
    230                N,
    231            );
    232            ptr::drop_in_place(initialized_part);
    233        }
    234    }
    235}
    236
    237#[unstable(feature = "iter_map_windows", issue = "87155")]
    238impl<I, F, R, const N: usize> Iterator for MapWindows<I, F, N>
    239where
    240    I: Iterator,
    241    F: FnMut(&[I::Item; N]) -> R,
    242{
    243    type Item = R;
    244
    245    fn next(&mut self) -> Option<Self::Item> {
    246        let window = self.inner.next_window()?;
    247        let out = (self.f)(window);
    248        Some(out)
    249    }
    250
    251    fn size_hint(&self) -> (usize, Option<usize>) {
    252        self.inner.size_hint()
    253    }
    254}
    255
    256// Note that even if the inner iterator not fused, the `MapWindows` is still fused,
    257// because we don't allow "holes" in the mapping window.
    258#[unstable(feature = "iter_map_windows", issue = "87155")]
    259impl<I, F, R, const N: usize> FusedIterator for MapWindows<I, F, N>
    260where
    261    I: Iterator,
    262    F: FnMut(&[I::Item; N]) -> R,
    263{
    264}
    265
    266#[unstable(feature = "iter_map_windows", issue = "87155")]
    267impl<I, F, R, const N: usize> ExactSizeIterator for MapWindows<I, F, N>
    268where
    269    I: ExactSizeIterator,
    270    F: FnMut(&[I::Item; N]) -> R,
    271{
    272}
    273
    274#[unstable(feature = "iter_map_windows", issue = "87155")]
    275impl<I: Iterator + fmt::Debug, F, const N: usize> fmt::Debug for MapWindows<I, F, N> {
    276    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    277        f.debug_struct("MapWindows").field("iter", &self.inner.iter).finish()
    278    }
    279}
    280
    281#[unstable(feature = "iter_map_windows", issue = "87155")]
    282impl<I, F, const N: usize> Clone for MapWindows<I, F, N>
    283where
    284    I: Iterator + Clone,
    285    F: Clone,
    286    I::Item: Clone,
    287{
    288    fn clone(&self) -> Self {
    289        Self { f: self.f.clone(), inner: self.inner.clone() }
    290    }
    291}