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", 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 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 = 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(to_drop: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: &[::Item; N] = self.inner.next_window()?;
247 let out: R = (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(name:"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}
292