1 | use crate::iter::adapters::SourceIter; |
2 | use crate::iter::{ |
3 | Cloned, Copied, Filter, FilterMap, Fuse, FusedIterator, InPlaceIterable, Map, TrustedFused, |
4 | TrustedLen, |
5 | }; |
6 | use crate::iter::{Empty, Once, OnceWith}; |
7 | use crate::num::NonZero; |
8 | use crate::ops::{ControlFlow, Try}; |
9 | use crate::result; |
10 | use crate::{array, fmt, option}; |
11 | |
12 | /// An iterator that maps each element to an iterator, and yields the elements |
13 | /// of the produced iterators. |
14 | /// |
15 | /// This `struct` is created by [`Iterator::flat_map`]. See its documentation |
16 | /// for more. |
17 | #[must_use = "iterators are lazy and do nothing unless consumed" ] |
18 | #[stable (feature = "rust1" , since = "1.0.0" )] |
19 | pub struct FlatMap<I, U: IntoIterator, F> { |
20 | inner: FlattenCompat<Map<I, F>, <U as IntoIterator>::IntoIter>, |
21 | } |
22 | |
23 | impl<I: Iterator, U: IntoIterator, F: FnMut(I::Item) -> U> FlatMap<I, U, F> { |
24 | pub(in crate::iter) fn new(iter: I, f: F) -> FlatMap<I, U, F> { |
25 | FlatMap { inner: FlattenCompat::new(iter:iter.map(f)) } |
26 | } |
27 | |
28 | pub(crate) fn into_parts(self) -> (Option<U::IntoIter>, Option<I>, Option<U::IntoIter>) { |
29 | ( |
30 | self.inner.frontiter, |
31 | self.inner.iter.into_inner().map(Map::into_inner), |
32 | self.inner.backiter, |
33 | ) |
34 | } |
35 | } |
36 | |
37 | #[stable (feature = "rust1" , since = "1.0.0" )] |
38 | impl<I: Clone, U, F: Clone> Clone for FlatMap<I, U, F> |
39 | where |
40 | U: Clone + IntoIterator<IntoIter: Clone>, |
41 | { |
42 | fn clone(&self) -> Self { |
43 | FlatMap { inner: self.inner.clone() } |
44 | } |
45 | } |
46 | |
47 | #[stable (feature = "core_impl_debug" , since = "1.9.0" )] |
48 | impl<I: fmt::Debug, U, F> fmt::Debug for FlatMap<I, U, F> |
49 | where |
50 | U: IntoIterator<IntoIter: fmt::Debug>, |
51 | { |
52 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
53 | f.debug_struct("FlatMap" ).field(name:"inner" , &self.inner).finish() |
54 | } |
55 | } |
56 | |
57 | #[stable (feature = "rust1" , since = "1.0.0" )] |
58 | impl<I: Iterator, U: IntoIterator, F> Iterator for FlatMap<I, U, F> |
59 | where |
60 | F: FnMut(I::Item) -> U, |
61 | { |
62 | type Item = U::Item; |
63 | |
64 | #[inline ] |
65 | fn next(&mut self) -> Option<U::Item> { |
66 | self.inner.next() |
67 | } |
68 | |
69 | #[inline ] |
70 | fn size_hint(&self) -> (usize, Option<usize>) { |
71 | self.inner.size_hint() |
72 | } |
73 | |
74 | #[inline ] |
75 | fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R |
76 | where |
77 | Self: Sized, |
78 | Fold: FnMut(Acc, Self::Item) -> R, |
79 | R: Try<Output = Acc>, |
80 | { |
81 | self.inner.try_fold(init, fold) |
82 | } |
83 | |
84 | #[inline ] |
85 | fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc |
86 | where |
87 | Fold: FnMut(Acc, Self::Item) -> Acc, |
88 | { |
89 | self.inner.fold(init, fold) |
90 | } |
91 | |
92 | #[inline ] |
93 | fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> { |
94 | self.inner.advance_by(n) |
95 | } |
96 | |
97 | #[inline ] |
98 | fn count(self) -> usize { |
99 | self.inner.count() |
100 | } |
101 | |
102 | #[inline ] |
103 | fn last(self) -> Option<Self::Item> { |
104 | self.inner.last() |
105 | } |
106 | } |
107 | |
108 | #[stable (feature = "rust1" , since = "1.0.0" )] |
109 | impl<I: DoubleEndedIterator, U, F> DoubleEndedIterator for FlatMap<I, U, F> |
110 | where |
111 | F: FnMut(I::Item) -> U, |
112 | U: IntoIterator<IntoIter: DoubleEndedIterator>, |
113 | { |
114 | #[inline ] |
115 | fn next_back(&mut self) -> Option<U::Item> { |
116 | self.inner.next_back() |
117 | } |
118 | |
119 | #[inline ] |
120 | fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R |
121 | where |
122 | Self: Sized, |
123 | Fold: FnMut(Acc, Self::Item) -> R, |
124 | R: Try<Output = Acc>, |
125 | { |
126 | self.inner.try_rfold(init, fold) |
127 | } |
128 | |
129 | #[inline ] |
130 | fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc |
131 | where |
132 | Fold: FnMut(Acc, Self::Item) -> Acc, |
133 | { |
134 | self.inner.rfold(init, fold) |
135 | } |
136 | |
137 | #[inline ] |
138 | fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> { |
139 | self.inner.advance_back_by(n) |
140 | } |
141 | } |
142 | |
143 | #[stable (feature = "fused" , since = "1.26.0" )] |
144 | impl<I, U, F> FusedIterator for FlatMap<I, U, F> |
145 | where |
146 | I: FusedIterator, |
147 | U: IntoIterator, |
148 | F: FnMut(I::Item) -> U, |
149 | { |
150 | } |
151 | |
152 | #[unstable (feature = "trusted_len" , issue = "37572" )] |
153 | unsafe impl<I, U, F> TrustedLen for FlatMap<I, U, F> |
154 | where |
155 | I: Iterator, |
156 | U: IntoIterator, |
157 | F: FnMut(I::Item) -> U, |
158 | FlattenCompat<Map<I, F>, <U as IntoIterator>::IntoIter>: TrustedLen, |
159 | { |
160 | } |
161 | |
162 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
163 | unsafe impl<I, U, F> InPlaceIterable for FlatMap<I, U, F> |
164 | where |
165 | I: InPlaceIterable, |
166 | U: BoundedSize + IntoIterator, |
167 | { |
168 | const EXPAND_BY: Option<NonZero<usize>> = const { |
169 | match (I::EXPAND_BY, U::UPPER_BOUND) { |
170 | (Some(m: NonZero), Some(n: NonZero)) => m.checked_mul(n), |
171 | _ => None, |
172 | } |
173 | }; |
174 | const MERGE_BY: Option<NonZero<usize>> = I::MERGE_BY; |
175 | } |
176 | |
177 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
178 | unsafe impl<I, U, F> SourceIter for FlatMap<I, U, F> |
179 | where |
180 | I: SourceIter + TrustedFused, |
181 | U: IntoIterator, |
182 | { |
183 | type Source = I::Source; |
184 | |
185 | #[inline ] |
186 | unsafe fn as_inner(&mut self) -> &mut I::Source { |
187 | // SAFETY: unsafe function forwarding to unsafe function with the same requirements |
188 | unsafe { SourceIter::as_inner(&mut self.inner.iter) } |
189 | } |
190 | } |
191 | |
192 | /// Marker trait for iterators/iterables which have a statically known upper |
193 | /// bound of the number of items they can produce. |
194 | /// |
195 | /// # Safety |
196 | /// |
197 | /// Implementations must not yield more elements than indicated by UPPER_BOUND if it is `Some`. |
198 | /// Used in specializations. Implementations must not be conditional on lifetimes or |
199 | /// user-implementable traits. |
200 | #[rustc_specialization_trait ] |
201 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
202 | unsafe trait BoundedSize { |
203 | const UPPER_BOUND: Option<NonZero<usize>> = NonZero::new(1); |
204 | } |
205 | |
206 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
207 | unsafe impl<T> BoundedSize for Option<T> {} |
208 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
209 | unsafe impl<T> BoundedSize for option::IntoIter<T> {} |
210 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
211 | unsafe impl<T, U> BoundedSize for Result<T, U> {} |
212 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
213 | unsafe impl<T> BoundedSize for result::IntoIter<T> {} |
214 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
215 | unsafe impl<T> BoundedSize for Once<T> {} |
216 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
217 | unsafe impl<T> BoundedSize for OnceWith<T> {} |
218 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
219 | unsafe impl<T, const N: usize> BoundedSize for [T; N] { |
220 | const UPPER_BOUND: Option<NonZero<usize>> = NonZero::new(N); |
221 | } |
222 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
223 | unsafe impl<T, const N: usize> BoundedSize for array::IntoIter<T, N> { |
224 | const UPPER_BOUND: Option<NonZero<usize>> = NonZero::new(N); |
225 | } |
226 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
227 | unsafe impl<I: BoundedSize, P> BoundedSize for Filter<I, P> { |
228 | const UPPER_BOUND: Option<NonZero<usize>> = I::UPPER_BOUND; |
229 | } |
230 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
231 | unsafe impl<I: BoundedSize, P> BoundedSize for FilterMap<I, P> { |
232 | const UPPER_BOUND: Option<NonZero<usize>> = I::UPPER_BOUND; |
233 | } |
234 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
235 | unsafe impl<I: BoundedSize, F> BoundedSize for Map<I, F> { |
236 | const UPPER_BOUND: Option<NonZero<usize>> = I::UPPER_BOUND; |
237 | } |
238 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
239 | unsafe impl<I: BoundedSize> BoundedSize for Copied<I> { |
240 | const UPPER_BOUND: Option<NonZero<usize>> = I::UPPER_BOUND; |
241 | } |
242 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
243 | unsafe impl<I: BoundedSize> BoundedSize for Cloned<I> { |
244 | const UPPER_BOUND: Option<NonZero<usize>> = I::UPPER_BOUND; |
245 | } |
246 | |
247 | /// An iterator that flattens one level of nesting in an iterator of things |
248 | /// that can be turned into iterators. |
249 | /// |
250 | /// This `struct` is created by the [`flatten`] method on [`Iterator`]. See its |
251 | /// documentation for more. |
252 | /// |
253 | /// [`flatten`]: Iterator::flatten() |
254 | #[must_use = "iterators are lazy and do nothing unless consumed" ] |
255 | #[stable (feature = "iterator_flatten" , since = "1.29.0" )] |
256 | pub struct Flatten<I: Iterator<Item: IntoIterator>> { |
257 | inner: FlattenCompat<I, <I::Item as IntoIterator>::IntoIter>, |
258 | } |
259 | |
260 | impl<I: Iterator<Item: IntoIterator>> Flatten<I> { |
261 | pub(in super::super) fn new(iter: I) -> Flatten<I> { |
262 | Flatten { inner: FlattenCompat::new(iter) } |
263 | } |
264 | } |
265 | |
266 | #[stable (feature = "iterator_flatten" , since = "1.29.0" )] |
267 | impl<I, U> fmt::Debug for Flatten<I> |
268 | where |
269 | I: fmt::Debug + Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>, |
270 | U: fmt::Debug + Iterator, |
271 | { |
272 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
273 | f.debug_struct("Flatten" ).field(name:"inner" , &self.inner).finish() |
274 | } |
275 | } |
276 | |
277 | #[stable (feature = "iterator_flatten" , since = "1.29.0" )] |
278 | impl<I, U> Clone for Flatten<I> |
279 | where |
280 | I: Clone + Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>, |
281 | U: Clone + Iterator, |
282 | { |
283 | fn clone(&self) -> Self { |
284 | Flatten { inner: self.inner.clone() } |
285 | } |
286 | } |
287 | |
288 | #[stable (feature = "iterator_flatten" , since = "1.29.0" )] |
289 | impl<I, U> Iterator for Flatten<I> |
290 | where |
291 | I: Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>, |
292 | U: Iterator, |
293 | { |
294 | type Item = U::Item; |
295 | |
296 | #[inline ] |
297 | fn next(&mut self) -> Option<U::Item> { |
298 | self.inner.next() |
299 | } |
300 | |
301 | #[inline ] |
302 | fn size_hint(&self) -> (usize, Option<usize>) { |
303 | self.inner.size_hint() |
304 | } |
305 | |
306 | #[inline ] |
307 | fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R |
308 | where |
309 | Self: Sized, |
310 | Fold: FnMut(Acc, Self::Item) -> R, |
311 | R: Try<Output = Acc>, |
312 | { |
313 | self.inner.try_fold(init, fold) |
314 | } |
315 | |
316 | #[inline ] |
317 | fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc |
318 | where |
319 | Fold: FnMut(Acc, Self::Item) -> Acc, |
320 | { |
321 | self.inner.fold(init, fold) |
322 | } |
323 | |
324 | #[inline ] |
325 | fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> { |
326 | self.inner.advance_by(n) |
327 | } |
328 | |
329 | #[inline ] |
330 | fn count(self) -> usize { |
331 | self.inner.count() |
332 | } |
333 | |
334 | #[inline ] |
335 | fn last(self) -> Option<Self::Item> { |
336 | self.inner.last() |
337 | } |
338 | } |
339 | |
340 | #[stable (feature = "iterator_flatten" , since = "1.29.0" )] |
341 | impl<I, U> DoubleEndedIterator for Flatten<I> |
342 | where |
343 | I: DoubleEndedIterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>, |
344 | U: DoubleEndedIterator, |
345 | { |
346 | #[inline ] |
347 | fn next_back(&mut self) -> Option<U::Item> { |
348 | self.inner.next_back() |
349 | } |
350 | |
351 | #[inline ] |
352 | fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R |
353 | where |
354 | Self: Sized, |
355 | Fold: FnMut(Acc, Self::Item) -> R, |
356 | R: Try<Output = Acc>, |
357 | { |
358 | self.inner.try_rfold(init, fold) |
359 | } |
360 | |
361 | #[inline ] |
362 | fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc |
363 | where |
364 | Fold: FnMut(Acc, Self::Item) -> Acc, |
365 | { |
366 | self.inner.rfold(init, fold) |
367 | } |
368 | |
369 | #[inline ] |
370 | fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> { |
371 | self.inner.advance_back_by(n) |
372 | } |
373 | } |
374 | |
375 | #[stable (feature = "iterator_flatten" , since = "1.29.0" )] |
376 | impl<I, U> FusedIterator for Flatten<I> |
377 | where |
378 | I: FusedIterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>, |
379 | U: Iterator, |
380 | { |
381 | } |
382 | |
383 | #[unstable (feature = "trusted_len" , issue = "37572" )] |
384 | unsafe impl<I> TrustedLen for Flatten<I> |
385 | where |
386 | I: Iterator<Item: IntoIterator>, |
387 | FlattenCompat<I, <I::Item as IntoIterator>::IntoIter>: TrustedLen, |
388 | { |
389 | } |
390 | |
391 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
392 | unsafe impl<I> InPlaceIterable for Flatten<I> |
393 | where |
394 | I: InPlaceIterable + Iterator, |
395 | <I as Iterator>::Item: IntoIterator + BoundedSize, |
396 | { |
397 | const EXPAND_BY: Option<NonZero<usize>> = const { |
398 | match (I::EXPAND_BY, I::Item::UPPER_BOUND) { |
399 | (Some(m: NonZero), Some(n: NonZero)) => m.checked_mul(n), |
400 | _ => None, |
401 | } |
402 | }; |
403 | const MERGE_BY: Option<NonZero<usize>> = I::MERGE_BY; |
404 | } |
405 | |
406 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
407 | unsafe impl<I> SourceIter for Flatten<I> |
408 | where |
409 | I: SourceIter + TrustedFused + Iterator, |
410 | <I as Iterator>::Item: IntoIterator, |
411 | { |
412 | type Source = I::Source; |
413 | |
414 | #[inline ] |
415 | unsafe fn as_inner(&mut self) -> &mut I::Source { |
416 | // SAFETY: unsafe function forwarding to unsafe function with the same requirements |
417 | unsafe { SourceIter::as_inner(&mut self.inner.iter) } |
418 | } |
419 | } |
420 | |
421 | #[stable (feature = "default_iters" , since = "1.70.0" )] |
422 | impl<I> Default for Flatten<I> |
423 | where |
424 | I: Default + Iterator<Item: IntoIterator>, |
425 | { |
426 | /// Creates a `Flatten` iterator from the default value of `I`. |
427 | /// |
428 | /// ``` |
429 | /// # use core::slice; |
430 | /// # use std::iter::Flatten; |
431 | /// let iter: Flatten<slice::Iter<'_, [u8; 4]>> = Default::default(); |
432 | /// assert_eq!(iter.count(), 0); |
433 | /// ``` |
434 | fn default() -> Self { |
435 | Flatten::new(iter:Default::default()) |
436 | } |
437 | } |
438 | |
439 | /// Real logic of both `Flatten` and `FlatMap` which simply delegate to |
440 | /// this type. |
441 | #[derive (Clone, Debug)] |
442 | #[unstable (feature = "trusted_len" , issue = "37572" )] |
443 | struct FlattenCompat<I, U> { |
444 | iter: Fuse<I>, |
445 | frontiter: Option<U>, |
446 | backiter: Option<U>, |
447 | } |
448 | impl<I, U> FlattenCompat<I, U> |
449 | where |
450 | I: Iterator, |
451 | { |
452 | /// Adapts an iterator by flattening it, for use in `flatten()` and `flat_map()`. |
453 | fn new(iter: I) -> FlattenCompat<I, U> { |
454 | FlattenCompat { iter: iter.fuse(), frontiter: None, backiter: None } |
455 | } |
456 | } |
457 | |
458 | impl<I, U> FlattenCompat<I, U> |
459 | where |
460 | I: Iterator<Item: IntoIterator<IntoIter = U>>, |
461 | { |
462 | /// Folds the inner iterators into an accumulator by applying an operation. |
463 | /// |
464 | /// Folds over the inner iterators, not over their elements. Is used by the `fold`, `count`, |
465 | /// and `last` methods. |
466 | #[inline ] |
467 | fn iter_fold<Acc, Fold>(self, mut acc: Acc, mut fold: Fold) -> Acc |
468 | where |
469 | Fold: FnMut(Acc, U) -> Acc, |
470 | { |
471 | #[inline ] |
472 | fn flatten<T: IntoIterator, Acc>( |
473 | fold: &mut impl FnMut(Acc, T::IntoIter) -> Acc, |
474 | ) -> impl FnMut(Acc, T) -> Acc + '_ { |
475 | move |acc, iter| fold(acc, iter.into_iter()) |
476 | } |
477 | |
478 | if let Some(iter) = self.frontiter { |
479 | acc = fold(acc, iter); |
480 | } |
481 | |
482 | acc = self.iter.fold(acc, flatten(&mut fold)); |
483 | |
484 | if let Some(iter) = self.backiter { |
485 | acc = fold(acc, iter); |
486 | } |
487 | |
488 | acc |
489 | } |
490 | |
491 | /// Folds over the inner iterators as long as the given function returns successfully, |
492 | /// always storing the most recent inner iterator in `self.frontiter`. |
493 | /// |
494 | /// Folds over the inner iterators, not over their elements. Is used by the `try_fold` and |
495 | /// `advance_by` methods. |
496 | #[inline ] |
497 | fn iter_try_fold<Acc, Fold, R>(&mut self, mut acc: Acc, mut fold: Fold) -> R |
498 | where |
499 | Fold: FnMut(Acc, &mut U) -> R, |
500 | R: Try<Output = Acc>, |
501 | { |
502 | #[inline ] |
503 | fn flatten<'a, T: IntoIterator, Acc, R: Try<Output = Acc>>( |
504 | frontiter: &'a mut Option<T::IntoIter>, |
505 | fold: &'a mut impl FnMut(Acc, &mut T::IntoIter) -> R, |
506 | ) -> impl FnMut(Acc, T) -> R + 'a { |
507 | move |acc, iter| fold(acc, frontiter.insert(iter.into_iter())) |
508 | } |
509 | |
510 | if let Some(iter) = &mut self.frontiter { |
511 | acc = fold(acc, iter)?; |
512 | } |
513 | self.frontiter = None; |
514 | |
515 | acc = self.iter.try_fold(acc, flatten(&mut self.frontiter, &mut fold))?; |
516 | self.frontiter = None; |
517 | |
518 | if let Some(iter) = &mut self.backiter { |
519 | acc = fold(acc, iter)?; |
520 | } |
521 | self.backiter = None; |
522 | |
523 | try { acc } |
524 | } |
525 | } |
526 | |
527 | impl<I, U> FlattenCompat<I, U> |
528 | where |
529 | I: DoubleEndedIterator<Item: IntoIterator<IntoIter = U>>, |
530 | { |
531 | /// Folds the inner iterators into an accumulator by applying an operation, starting form the |
532 | /// back. |
533 | /// |
534 | /// Folds over the inner iterators, not over their elements. Is used by the `rfold` method. |
535 | #[inline ] |
536 | fn iter_rfold<Acc, Fold>(self, mut acc: Acc, mut fold: Fold) -> Acc |
537 | where |
538 | Fold: FnMut(Acc, U) -> Acc, |
539 | { |
540 | #[inline ] |
541 | fn flatten<T: IntoIterator, Acc>( |
542 | fold: &mut impl FnMut(Acc, T::IntoIter) -> Acc, |
543 | ) -> impl FnMut(Acc, T) -> Acc + '_ { |
544 | move |acc, iter| fold(acc, iter.into_iter()) |
545 | } |
546 | |
547 | if let Some(iter) = self.backiter { |
548 | acc = fold(acc, iter); |
549 | } |
550 | |
551 | acc = self.iter.rfold(acc, flatten(&mut fold)); |
552 | |
553 | if let Some(iter) = self.frontiter { |
554 | acc = fold(acc, iter); |
555 | } |
556 | |
557 | acc |
558 | } |
559 | |
560 | /// Folds over the inner iterators in reverse order as long as the given function returns |
561 | /// successfully, always storing the most recent inner iterator in `self.backiter`. |
562 | /// |
563 | /// Folds over the inner iterators, not over their elements. Is used by the `try_rfold` and |
564 | /// `advance_back_by` methods. |
565 | #[inline ] |
566 | fn iter_try_rfold<Acc, Fold, R>(&mut self, mut acc: Acc, mut fold: Fold) -> R |
567 | where |
568 | Fold: FnMut(Acc, &mut U) -> R, |
569 | R: Try<Output = Acc>, |
570 | { |
571 | #[inline ] |
572 | fn flatten<'a, T: IntoIterator, Acc, R: Try>( |
573 | backiter: &'a mut Option<T::IntoIter>, |
574 | fold: &'a mut impl FnMut(Acc, &mut T::IntoIter) -> R, |
575 | ) -> impl FnMut(Acc, T) -> R + 'a { |
576 | move |acc, iter| fold(acc, backiter.insert(iter.into_iter())) |
577 | } |
578 | |
579 | if let Some(iter) = &mut self.backiter { |
580 | acc = fold(acc, iter)?; |
581 | } |
582 | self.backiter = None; |
583 | |
584 | acc = self.iter.try_rfold(acc, flatten(&mut self.backiter, &mut fold))?; |
585 | self.backiter = None; |
586 | |
587 | if let Some(iter) = &mut self.frontiter { |
588 | acc = fold(acc, iter)?; |
589 | } |
590 | self.frontiter = None; |
591 | |
592 | try { acc } |
593 | } |
594 | } |
595 | |
596 | // See also the `OneShot` specialization below. |
597 | impl<I, U> Iterator for FlattenCompat<I, U> |
598 | where |
599 | I: Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>, |
600 | U: Iterator, |
601 | { |
602 | type Item = U::Item; |
603 | |
604 | #[inline ] |
605 | default fn next(&mut self) -> Option<U::Item> { |
606 | loop { |
607 | if let elt @ Some(_) = and_then_or_clear(&mut self.frontiter, Iterator::next) { |
608 | return elt; |
609 | } |
610 | match self.iter.next() { |
611 | None => return and_then_or_clear(&mut self.backiter, Iterator::next), |
612 | Some(inner) => self.frontiter = Some(inner.into_iter()), |
613 | } |
614 | } |
615 | } |
616 | |
617 | #[inline ] |
618 | default fn size_hint(&self) -> (usize, Option<usize>) { |
619 | let (flo, fhi) = self.frontiter.as_ref().map_or((0, Some(0)), U::size_hint); |
620 | let (blo, bhi) = self.backiter.as_ref().map_or((0, Some(0)), U::size_hint); |
621 | let lo = flo.saturating_add(blo); |
622 | |
623 | if let Some(fixed_size) = <<I as Iterator>::Item as ConstSizeIntoIterator>::size() { |
624 | let (lower, upper) = self.iter.size_hint(); |
625 | |
626 | let lower = lower.saturating_mul(fixed_size).saturating_add(lo); |
627 | let upper = |
628 | try { fhi?.checked_add(bhi?)?.checked_add(fixed_size.checked_mul(upper?)?)? }; |
629 | |
630 | return (lower, upper); |
631 | } |
632 | |
633 | match (self.iter.size_hint(), fhi, bhi) { |
634 | ((0, Some(0)), Some(a), Some(b)) => (lo, a.checked_add(b)), |
635 | _ => (lo, None), |
636 | } |
637 | } |
638 | |
639 | #[inline ] |
640 | default fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R |
641 | where |
642 | Self: Sized, |
643 | Fold: FnMut(Acc, Self::Item) -> R, |
644 | R: Try<Output = Acc>, |
645 | { |
646 | #[inline ] |
647 | fn flatten<U: Iterator, Acc, R: Try<Output = Acc>>( |
648 | mut fold: impl FnMut(Acc, U::Item) -> R, |
649 | ) -> impl FnMut(Acc, &mut U) -> R { |
650 | move |acc, iter| iter.try_fold(acc, &mut fold) |
651 | } |
652 | |
653 | self.iter_try_fold(init, flatten(fold)) |
654 | } |
655 | |
656 | #[inline ] |
657 | default fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc |
658 | where |
659 | Fold: FnMut(Acc, Self::Item) -> Acc, |
660 | { |
661 | #[inline ] |
662 | fn flatten<U: Iterator, Acc>( |
663 | mut fold: impl FnMut(Acc, U::Item) -> Acc, |
664 | ) -> impl FnMut(Acc, U) -> Acc { |
665 | move |acc, iter| iter.fold(acc, &mut fold) |
666 | } |
667 | |
668 | self.iter_fold(init, flatten(fold)) |
669 | } |
670 | |
671 | #[inline ] |
672 | #[rustc_inherit_overflow_checks ] |
673 | default fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> { |
674 | #[inline ] |
675 | #[rustc_inherit_overflow_checks ] |
676 | fn advance<U: Iterator>(n: usize, iter: &mut U) -> ControlFlow<(), usize> { |
677 | match iter.advance_by(n) { |
678 | Ok(()) => ControlFlow::Break(()), |
679 | Err(remaining) => ControlFlow::Continue(remaining.get()), |
680 | } |
681 | } |
682 | |
683 | match self.iter_try_fold(n, advance) { |
684 | ControlFlow::Continue(remaining) => NonZero::new(remaining).map_or(Ok(()), Err), |
685 | _ => Ok(()), |
686 | } |
687 | } |
688 | |
689 | #[inline ] |
690 | default fn count(self) -> usize { |
691 | #[inline ] |
692 | #[rustc_inherit_overflow_checks ] |
693 | fn count<U: Iterator>(acc: usize, iter: U) -> usize { |
694 | acc + iter.count() |
695 | } |
696 | |
697 | self.iter_fold(0, count) |
698 | } |
699 | |
700 | #[inline ] |
701 | default fn last(self) -> Option<Self::Item> { |
702 | #[inline ] |
703 | fn last<U: Iterator>(last: Option<U::Item>, iter: U) -> Option<U::Item> { |
704 | iter.last().or(last) |
705 | } |
706 | |
707 | self.iter_fold(None, last) |
708 | } |
709 | } |
710 | |
711 | // See also the `OneShot` specialization below. |
712 | impl<I, U> DoubleEndedIterator for FlattenCompat<I, U> |
713 | where |
714 | I: DoubleEndedIterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>, |
715 | U: DoubleEndedIterator, |
716 | { |
717 | #[inline ] |
718 | default fn next_back(&mut self) -> Option<U::Item> { |
719 | loop { |
720 | if let elt @ Some(_) = and_then_or_clear(&mut self.backiter, |b| b.next_back()) { |
721 | return elt; |
722 | } |
723 | match self.iter.next_back() { |
724 | None => return and_then_or_clear(&mut self.frontiter, |f| f.next_back()), |
725 | Some(inner) => self.backiter = Some(inner.into_iter()), |
726 | } |
727 | } |
728 | } |
729 | |
730 | #[inline ] |
731 | default fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R |
732 | where |
733 | Self: Sized, |
734 | Fold: FnMut(Acc, Self::Item) -> R, |
735 | R: Try<Output = Acc>, |
736 | { |
737 | #[inline ] |
738 | fn flatten<U: DoubleEndedIterator, Acc, R: Try<Output = Acc>>( |
739 | mut fold: impl FnMut(Acc, U::Item) -> R, |
740 | ) -> impl FnMut(Acc, &mut U) -> R { |
741 | move |acc, iter| iter.try_rfold(acc, &mut fold) |
742 | } |
743 | |
744 | self.iter_try_rfold(init, flatten(fold)) |
745 | } |
746 | |
747 | #[inline ] |
748 | default fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc |
749 | where |
750 | Fold: FnMut(Acc, Self::Item) -> Acc, |
751 | { |
752 | #[inline ] |
753 | fn flatten<U: DoubleEndedIterator, Acc>( |
754 | mut fold: impl FnMut(Acc, U::Item) -> Acc, |
755 | ) -> impl FnMut(Acc, U) -> Acc { |
756 | move |acc, iter| iter.rfold(acc, &mut fold) |
757 | } |
758 | |
759 | self.iter_rfold(init, flatten(fold)) |
760 | } |
761 | |
762 | #[inline ] |
763 | #[rustc_inherit_overflow_checks ] |
764 | default fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> { |
765 | #[inline ] |
766 | #[rustc_inherit_overflow_checks ] |
767 | fn advance<U: DoubleEndedIterator>(n: usize, iter: &mut U) -> ControlFlow<(), usize> { |
768 | match iter.advance_back_by(n) { |
769 | Ok(()) => ControlFlow::Break(()), |
770 | Err(remaining) => ControlFlow::Continue(remaining.get()), |
771 | } |
772 | } |
773 | |
774 | match self.iter_try_rfold(n, advance) { |
775 | ControlFlow::Continue(remaining) => NonZero::new(remaining).map_or(Ok(()), Err), |
776 | _ => Ok(()), |
777 | } |
778 | } |
779 | } |
780 | |
781 | unsafe impl<const N: usize, I, T> TrustedLen |
782 | for FlattenCompat<I, <[T; N] as IntoIterator>::IntoIter> |
783 | where |
784 | I: TrustedLen<Item = [T; N]>, |
785 | { |
786 | } |
787 | |
788 | unsafe impl<'a, const N: usize, I, T> TrustedLen |
789 | for FlattenCompat<I, <&'a [T; N] as IntoIterator>::IntoIter> |
790 | where |
791 | I: TrustedLen<Item = &'a [T; N]>, |
792 | { |
793 | } |
794 | |
795 | unsafe impl<'a, const N: usize, I, T> TrustedLen |
796 | for FlattenCompat<I, <&'a mut [T; N] as IntoIterator>::IntoIter> |
797 | where |
798 | I: TrustedLen<Item = &'a mut [T; N]>, |
799 | { |
800 | } |
801 | |
802 | trait ConstSizeIntoIterator: IntoIterator { |
803 | // FIXME(#31844): convert to an associated const once specialization supports that |
804 | fn size() -> Option<usize>; |
805 | } |
806 | |
807 | impl<T> ConstSizeIntoIterator for T |
808 | where |
809 | T: IntoIterator, |
810 | { |
811 | #[inline ] |
812 | default fn size() -> Option<usize> { |
813 | None |
814 | } |
815 | } |
816 | |
817 | impl<T, const N: usize> ConstSizeIntoIterator for [T; N] { |
818 | #[inline ] |
819 | fn size() -> Option<usize> { |
820 | Some(N) |
821 | } |
822 | } |
823 | |
824 | impl<T, const N: usize> ConstSizeIntoIterator for &[T; N] { |
825 | #[inline ] |
826 | fn size() -> Option<usize> { |
827 | Some(N) |
828 | } |
829 | } |
830 | |
831 | impl<T, const N: usize> ConstSizeIntoIterator for &mut [T; N] { |
832 | #[inline ] |
833 | fn size() -> Option<usize> { |
834 | Some(N) |
835 | } |
836 | } |
837 | |
838 | #[inline ] |
839 | fn and_then_or_clear<T, U>(opt: &mut Option<T>, f: impl FnOnce(&mut T) -> Option<U>) -> Option<U> { |
840 | let x: Option = f(opt.as_mut()?); |
841 | if x.is_none() { |
842 | *opt = None; |
843 | } |
844 | x |
845 | } |
846 | |
847 | /// Specialization trait for iterator types that never return more than one item. |
848 | /// |
849 | /// Note that we still have to deal with the possibility that the iterator was |
850 | /// already exhausted before it came into our control. |
851 | #[rustc_specialization_trait ] |
852 | trait OneShot {} |
853 | |
854 | // These all have exactly one item, if not already consumed. |
855 | impl<T> OneShot for Once<T> {} |
856 | impl<F> OneShot for OnceWith<F> {} |
857 | impl<T> OneShot for array::IntoIter<T, 1> {} |
858 | impl<T> OneShot for option::IntoIter<T> {} |
859 | impl<T> OneShot for option::Iter<'_, T> {} |
860 | impl<T> OneShot for option::IterMut<'_, T> {} |
861 | impl<T> OneShot for result::IntoIter<T> {} |
862 | impl<T> OneShot for result::Iter<'_, T> {} |
863 | impl<T> OneShot for result::IterMut<'_, T> {} |
864 | |
865 | // These are always empty, which is fine to optimize too. |
866 | impl<T> OneShot for Empty<T> {} |
867 | impl<T> OneShot for array::IntoIter<T, 0> {} |
868 | |
869 | // These adaptors never increase the number of items. |
870 | // (There are more possible, but for now this matches BoundedSize above.) |
871 | impl<I: OneShot> OneShot for Cloned<I> {} |
872 | impl<I: OneShot> OneShot for Copied<I> {} |
873 | impl<I: OneShot, P> OneShot for Filter<I, P> {} |
874 | impl<I: OneShot, P> OneShot for FilterMap<I, P> {} |
875 | impl<I: OneShot, F> OneShot for Map<I, F> {} |
876 | |
877 | // Blanket impls pass this property through as well |
878 | // (but we can't do `Box<I>` unless we expose this trait to alloc) |
879 | impl<I: OneShot> OneShot for &mut I {} |
880 | |
881 | #[inline ] |
882 | fn into_item<I>(inner: I) -> Option<I::Item> |
883 | where |
884 | I: IntoIterator<IntoIter: OneShot>, |
885 | { |
886 | inner.into_iter().next() |
887 | } |
888 | |
889 | #[inline ] |
890 | fn flatten_one<I: IntoIterator<IntoIter: OneShot>, Acc>( |
891 | mut fold: impl FnMut(Acc, I::Item) -> Acc, |
892 | ) -> impl FnMut(Acc, I) -> Acc { |
893 | move |acc: Acc, inner: I| match inner.into_iter().next() { |
894 | Some(item: ::Item) => fold(acc, item), |
895 | None => acc, |
896 | } |
897 | } |
898 | |
899 | #[inline ] |
900 | fn try_flatten_one<I: IntoIterator<IntoIter: OneShot>, Acc, R: Try<Output = Acc>>( |
901 | mut fold: impl FnMut(Acc, I::Item) -> R, |
902 | ) -> impl FnMut(Acc, I) -> R { |
903 | move |acc: Acc, inner: I| match inner.into_iter().next() { |
904 | Some(item: ::Item) => fold(acc, item), |
905 | None => try { acc }, |
906 | } |
907 | } |
908 | |
909 | #[inline ] |
910 | fn advance_by_one<I>(n: NonZero<usize>, inner: I) -> Option<NonZero<usize>> |
911 | where |
912 | I: IntoIterator<IntoIter: OneShot>, |
913 | { |
914 | match inner.into_iter().next() { |
915 | Some(_) => NonZero::new(n.get() - 1), |
916 | None => Some(n), |
917 | } |
918 | } |
919 | |
920 | // Specialization: When the inner iterator `U` never returns more than one item, the `frontiter` and |
921 | // `backiter` states are a waste, because they'll always have already consumed their item. So in |
922 | // this impl, we completely ignore them and just focus on `self.iter`, and we only call the inner |
923 | // `U::next()` one time. |
924 | // |
925 | // It's mostly fine if we accidentally mix this with the more generic impls, e.g. by forgetting to |
926 | // specialize one of the methods. If the other impl did set the front or back, we wouldn't see it |
927 | // here, but it would be empty anyway; and if the other impl looked for a front or back that we |
928 | // didn't bother setting, it would just see `None` (or a previous empty) and move on. |
929 | // |
930 | // An exception to that is `advance_by(0)` and `advance_back_by(0)`, where the generic impls may set |
931 | // `frontiter` or `backiter` without consuming the item, so we **must** override those. |
932 | impl<I, U> Iterator for FlattenCompat<I, U> |
933 | where |
934 | I: Iterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>, |
935 | U: Iterator + OneShot, |
936 | { |
937 | #[inline ] |
938 | fn next(&mut self) -> Option<U::Item> { |
939 | while let Some(inner) = self.iter.next() { |
940 | if let item @ Some(_) = inner.into_iter().next() { |
941 | return item; |
942 | } |
943 | } |
944 | None |
945 | } |
946 | |
947 | #[inline ] |
948 | fn size_hint(&self) -> (usize, Option<usize>) { |
949 | let (lower, upper) = self.iter.size_hint(); |
950 | match <I::Item as ConstSizeIntoIterator>::size() { |
951 | Some(0) => (0, Some(0)), |
952 | Some(1) => (lower, upper), |
953 | _ => (0, upper), |
954 | } |
955 | } |
956 | |
957 | #[inline ] |
958 | fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R |
959 | where |
960 | Self: Sized, |
961 | Fold: FnMut(Acc, Self::Item) -> R, |
962 | R: Try<Output = Acc>, |
963 | { |
964 | self.iter.try_fold(init, try_flatten_one(fold)) |
965 | } |
966 | |
967 | #[inline ] |
968 | fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc |
969 | where |
970 | Fold: FnMut(Acc, Self::Item) -> Acc, |
971 | { |
972 | self.iter.fold(init, flatten_one(fold)) |
973 | } |
974 | |
975 | #[inline ] |
976 | fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> { |
977 | if let Some(n) = NonZero::new(n) { |
978 | self.iter.try_fold(n, advance_by_one).map_or(Ok(()), Err) |
979 | } else { |
980 | // Just advance the outer iterator |
981 | self.iter.advance_by(0) |
982 | } |
983 | } |
984 | |
985 | #[inline ] |
986 | fn count(self) -> usize { |
987 | self.iter.filter_map(into_item).count() |
988 | } |
989 | |
990 | #[inline ] |
991 | fn last(self) -> Option<Self::Item> { |
992 | self.iter.filter_map(into_item).last() |
993 | } |
994 | } |
995 | |
996 | // Note: We don't actually care about `U: DoubleEndedIterator`, since forward and backward are the |
997 | // same for a one-shot iterator, but we have to keep that to match the default specialization. |
998 | impl<I, U> DoubleEndedIterator for FlattenCompat<I, U> |
999 | where |
1000 | I: DoubleEndedIterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>, |
1001 | U: DoubleEndedIterator + OneShot, |
1002 | { |
1003 | #[inline ] |
1004 | fn next_back(&mut self) -> Option<U::Item> { |
1005 | while let Some(inner) = self.iter.next_back() { |
1006 | if let item @ Some(_) = inner.into_iter().next() { |
1007 | return item; |
1008 | } |
1009 | } |
1010 | None |
1011 | } |
1012 | |
1013 | #[inline ] |
1014 | fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R |
1015 | where |
1016 | Self: Sized, |
1017 | Fold: FnMut(Acc, Self::Item) -> R, |
1018 | R: Try<Output = Acc>, |
1019 | { |
1020 | self.iter.try_rfold(init, try_flatten_one(fold)) |
1021 | } |
1022 | |
1023 | #[inline ] |
1024 | fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc |
1025 | where |
1026 | Fold: FnMut(Acc, Self::Item) -> Acc, |
1027 | { |
1028 | self.iter.rfold(init, flatten_one(fold)) |
1029 | } |
1030 | |
1031 | #[inline ] |
1032 | fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> { |
1033 | if let Some(n) = NonZero::new(n) { |
1034 | self.iter.try_rfold(n, advance_by_one).map_or(Ok(()), Err) |
1035 | } else { |
1036 | // Just advance the outer iterator |
1037 | self.iter.advance_back_by(0) |
1038 | } |
1039 | } |
1040 | } |
1041 | |