flatten.rs - source
Description: Source of the Rust file `library/core/src/iter/adapters/flatten.rs`.

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