core/iter/adapters/
map_windows.rs

1use crate::iter::FusedIterator;
2use crate::mem::{self, 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", reason = "recently added", 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 mem::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::<T>()) };
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", reason = "recently added", 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", reason = "recently added", 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", reason = "recently added", 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", reason = "recently added", 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", reason = "recently added", 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}