1use crate::iter::adapters::SourceIter;
2use crate::iter::{
3 Cloned, Copied, Filter, FilterMap, Fuse, FusedIterator, InPlaceIterable, Map, TrustedFused,
4 TrustedLen,
5};
6use crate::iter::{Empty, Once, OnceWith};
7use crate::num::NonZero;
8use crate::ops::{ControlFlow, Try};
9use crate::result;
10use 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")]
19pub struct FlatMap<I, U: IntoIterator, F> {
20 inner: FlattenCompat<Map<I, F>, <U as IntoIterator>::IntoIter>,
21}
22
23impl<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")]
38impl<I: Clone, U, F: Clone> Clone for FlatMap<I, U, F>
39where
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")]
48impl<I: fmt::Debug, U, F> fmt::Debug for FlatMap<I, U, F>
49where
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")]
58impl<I: Iterator, U: IntoIterator, F> Iterator for FlatMap<I, U, F>
59where
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")]
109impl<I: DoubleEndedIterator, U, F> DoubleEndedIterator for FlatMap<I, U, F>
110where
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")]
144impl<I, U, F> FusedIterator for FlatMap<I, U, F>
145where
146 I: FusedIterator,
147 U: IntoIterator,
148 F: FnMut(I::Item) -> U,
149{
150}
151
152#[unstable(feature = "trusted_len", issue = "37572")]
153unsafe impl<I, U, F> TrustedLen for FlatMap<I, U, F>
154where
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")]
163unsafe impl<I, U, F> InPlaceIterable for FlatMap<I, U, F>
164where
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")]
178unsafe impl<I, U, F> SourceIter for FlatMap<I, U, F>
179where
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")]
202unsafe trait BoundedSize {
203 const UPPER_BOUND: Option<NonZero<usize>> = NonZero::new(1);
204}
205
206#[unstable(issue = "none", feature = "inplace_iteration")]
207unsafe impl<T> BoundedSize for Option<T> {}
208#[unstable(issue = "none", feature = "inplace_iteration")]
209unsafe impl<T> BoundedSize for option::IntoIter<T> {}
210#[unstable(issue = "none", feature = "inplace_iteration")]
211unsafe impl<T, U> BoundedSize for Result<T, U> {}
212#[unstable(issue = "none", feature = "inplace_iteration")]
213unsafe impl<T> BoundedSize for result::IntoIter<T> {}
214#[unstable(issue = "none", feature = "inplace_iteration")]
215unsafe impl<T> BoundedSize for Once<T> {}
216#[unstable(issue = "none", feature = "inplace_iteration")]
217unsafe impl<T> BoundedSize for OnceWith<T> {}
218#[unstable(issue = "none", feature = "inplace_iteration")]
219unsafe 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")]
223unsafe 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")]
227unsafe 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")]
231unsafe 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")]
235unsafe 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")]
239unsafe 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")]
243unsafe 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")]
256pub struct Flatten<I: Iterator<Item: IntoIterator>> {
257 inner: FlattenCompat<I, <I::Item as IntoIterator>::IntoIter>,
258}
259
260impl<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")]
267impl<I, U> fmt::Debug for Flatten<I>
268where
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")]
278impl<I, U> Clone for Flatten<I>
279where
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")]
289impl<I, U> Iterator for Flatten<I>
290where
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")]
341impl<I, U> DoubleEndedIterator for Flatten<I>
342where
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")]
376impl<I, U> FusedIterator for Flatten<I>
377where
378 I: FusedIterator<Item: IntoIterator<IntoIter = U, Item = U::Item>>,
379 U: Iterator,
380{
381}
382
383#[unstable(feature = "trusted_len", issue = "37572")]
384unsafe impl<I> TrustedLen for Flatten<I>
385where
386 I: Iterator<Item: IntoIterator>,
387 FlattenCompat<I, <I::Item as IntoIterator>::IntoIter>: TrustedLen,
388{
389}
390
391#[unstable(issue = "none", feature = "inplace_iteration")]
392unsafe impl<I> InPlaceIterable for Flatten<I>
393where
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")]
407unsafe impl<I> SourceIter for Flatten<I>
408where
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")]
422impl<I> Default for Flatten<I>
423where
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")]
443struct FlattenCompat<I, U> {
444 iter: Fuse<I>,
445 frontiter: Option<U>,
446 backiter: Option<U>,
447}
448impl<I, U> FlattenCompat<I, U>
449where
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
458impl<I, U> FlattenCompat<I, U>
459where
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
527impl<I, U> FlattenCompat<I, U>
528where
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.
597impl<I, U> Iterator for FlattenCompat<I, U>
598where
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.
712impl<I, U> DoubleEndedIterator for FlattenCompat<I, U>
713where
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
781unsafe impl<const N: usize, I, T> TrustedLen
782 for FlattenCompat<I, <[T; N] as IntoIterator>::IntoIter>
783where
784 I: TrustedLen<Item = [T; N]>,
785{
786}
787
788unsafe impl<'a, const N: usize, I, T> TrustedLen
789 for FlattenCompat<I, <&'a [T; N] as IntoIterator>::IntoIter>
790where
791 I: TrustedLen<Item = &'a [T; N]>,
792{
793}
794
795unsafe impl<'a, const N: usize, I, T> TrustedLen
796 for FlattenCompat<I, <&'a mut [T; N] as IntoIterator>::IntoIter>
797where
798 I: TrustedLen<Item = &'a mut [T; N]>,
799{
800}
801
802trait ConstSizeIntoIterator: IntoIterator {
803 // FIXME(#31844): convert to an associated const once specialization supports that
804 fn size() -> Option<usize>;
805}
806
807impl<T> ConstSizeIntoIterator for T
808where
809 T: IntoIterator,
810{
811 #[inline]
812 default fn size() -> Option<usize> {
813 None
814 }
815}
816
817impl<T, const N: usize> ConstSizeIntoIterator for [T; N] {
818 #[inline]
819 fn size() -> Option<usize> {
820 Some(N)
821 }
822}
823
824impl<T, const N: usize> ConstSizeIntoIterator for &[T; N] {
825 #[inline]
826 fn size() -> Option<usize> {
827 Some(N)
828 }
829}
830
831impl<T, const N: usize> ConstSizeIntoIterator for &mut [T; N] {
832 #[inline]
833 fn size() -> Option<usize> {
834 Some(N)
835 }
836}
837
838#[inline]
839fn 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]
852trait OneShot {}
853
854// These all have exactly one item, if not already consumed.
855impl<T> OneShot for Once<T> {}
856impl<F> OneShot for OnceWith<F> {}
857impl<T> OneShot for array::IntoIter<T, 1> {}
858impl<T> OneShot for option::IntoIter<T> {}
859impl<T> OneShot for option::Iter<'_, T> {}
860impl<T> OneShot for option::IterMut<'_, T> {}
861impl<T> OneShot for result::IntoIter<T> {}
862impl<T> OneShot for result::Iter<'_, T> {}
863impl<T> OneShot for result::IterMut<'_, T> {}
864
865// These are always empty, which is fine to optimize too.
866impl<T> OneShot for Empty<T> {}
867impl<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.)
871impl<I: OneShot> OneShot for Cloned<I> {}
872impl<I: OneShot> OneShot for Copied<I> {}
873impl<I: OneShot, P> OneShot for Filter<I, P> {}
874impl<I: OneShot, P> OneShot for FilterMap<I, P> {}
875impl<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)
879impl<I: OneShot> OneShot for &mut I {}
880
881#[inline]
882fn into_item<I>(inner: I) -> Option<I::Item>
883where
884 I: IntoIterator<IntoIter: OneShot>,
885{
886 inner.into_iter().next()
887}
888
889#[inline]
890fn 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]
900fn 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]
910fn advance_by_one<I>(n: NonZero<usize>, inner: I) -> Option<NonZero<usize>>
911where
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.
932impl<I, U> Iterator for FlattenCompat<I, U>
933where
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.
998impl<I, U> DoubleEndedIterator for FlattenCompat<I, U>
999where
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