core/iter/adapters/
flatten.rs

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