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

core/iter/adapters/
map_windows.rs
