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

#
core/iter/adapters/
chain.rs

    1use crate::iter::{FusedIterator, TrustedLen};
    2use crate::num::NonZero;
    3use crate::ops::Try;
    4
    5/// An iterator that links two iterators together, in a chain.
    6///
    7/// This `struct` is created by [`chain`] or [`Iterator::chain`]. See their
    8/// documentation for more.
    9///
    10/// # Examples
    11///
    12/// ```
    13/// use std::iter::Chain;
    14/// use std::slice::Iter;
    15///
    16/// let a1 = [1, 2, 3];
    17/// let a2 = [4, 5, 6];
    18/// let iter: Chain<Iter<'_, _>, Iter<'_, _>> = a1.iter().chain(a2.iter());
    19/// ```
    20#[derive(Clone, Debug)]
    21#[must_use = "iterators are lazy and do nothing unless consumed"]
    22#[stable(feature = "rust1", since = "1.0.0")]
    23pub struct Chain<A, B> {
    24    // These are "fused" with `Option` so we don't need separate state to track which part is
    25    // already exhausted, and we may also get niche layout for `None`. We don't use the real `Fuse`
    26    // adapter because its specialization for `FusedIterator` unconditionally descends into the
    27    // iterator, and that could be expensive to keep revisiting stuff like nested chains. It also
    28    // hurts compiler performance to add more iterator layers to `Chain`.
    29    //
    30    // Only the "first" iterator is actually set `None` when exhausted, depending on whether you
    31    // iterate forward or backward. If you mix directions, then both sides may be `None`.
    32    a: Option<A>,
    33    b: Option<B>,
    34}
    35impl<A, B> Chain<A, B> {
    36    pub(in super::super) fn new(a: A, b: B) -> Chain<A, B> {
    37        Chain { a: Some(a), b: Some(b) }
    38    }
    39}
    40
    41/// Converts the arguments to iterators and links them together, in a chain.
    42///
    43/// See the documentation of [`Iterator::chain`] for more.
    44///
    45/// # Examples
    46///
    47/// ```
    48/// use std::iter::chain;
    49///
    50/// let a = [1, 2, 3];
    51/// let b = [4, 5, 6];
    52///
    53/// let mut iter = chain(a, b);
    54///
    55/// assert_eq!(iter.next(), Some(1));
    56/// assert_eq!(iter.next(), Some(2));
    57/// assert_eq!(iter.next(), Some(3));
    58/// assert_eq!(iter.next(), Some(4));
    59/// assert_eq!(iter.next(), Some(5));
    60/// assert_eq!(iter.next(), Some(6));
    61/// assert_eq!(iter.next(), None);
    62/// ```
    63#[stable(feature = "iter_chain", since = "1.91.0")]
    64pub fn chain<A, B>(a: A, b: B) -> Chain<A::IntoIter, B::IntoIter>
    65where
    66    A: IntoIterator,
    67    B: IntoIterator<Item = A::Item>,
    68{
    69    Chain::new(a.into_iter(), b.into_iter())
    70}
    71
    72#[stable(feature = "rust1", since = "1.0.0")]
    73impl<A, B> Iterator for Chain<A, B>
    74where
    75    A: Iterator,
    76    B: Iterator<Item = A::Item>,
    77{
    78    type Item = A::Item;
    79
    80    #[inline]
    81    fn next(&mut self) -> Option<A::Item> {
    82        and_then_or_clear(&mut self.a, Iterator::next).or_else(|| self.b.as_mut()?.next())
    83    }
    84
    85    #[inline]
    86    #[rustc_inherit_overflow_checks]
    87    fn count(self) -> usize {
    88        let a_count = match self.a {
    89            Some(a) => a.count(),
    90            None => 0,
    91        };
    92        let b_count = match self.b {
    93            Some(b) => b.count(),
    94            None => 0,
    95        };
    96        a_count + b_count
    97    }
    98
    99    fn try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R
    100    where
    101        Self: Sized,
    102        F: FnMut(Acc, Self::Item) -> R,
    103        R: Try<Output = Acc>,
    104    {
    105        if let Some(ref mut a) = self.a {
    106            acc = a.try_fold(acc, &mut f)?;
    107            self.a = None;
    108        }
    109        if let Some(ref mut b) = self.b {
    110            acc = b.try_fold(acc, f)?;
    111            // we don't fuse the second iterator
    112        }
    113        try { acc }
    114    }
    115
    116    fn fold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc
    117    where
    118        F: FnMut(Acc, Self::Item) -> Acc,
    119    {
    120        if let Some(a) = self.a {
    121            acc = a.fold(acc, &mut f);
    122        }
    123        if let Some(b) = self.b {
    124            acc = b.fold(acc, f);
    125        }
    126        acc
    127    }
    128
    129    #[inline]
    130    fn advance_by(&mut self, mut n: usize) -> Result<(), NonZero<usize>> {
    131        if let Some(ref mut a) = self.a {
    132            n = match a.advance_by(n) {
    133                Ok(()) => return Ok(()),
    134                Err(k) => k.get(),
    135            };
    136            self.a = None;
    137        }
    138
    139        if let Some(ref mut b) = self.b {
    140            return b.advance_by(n);
    141            // we don't fuse the second iterator
    142        }
    143
    144        NonZero::new(n).map_or(Ok(()), Err)
    145    }
    146
    147    #[inline]
    148    fn nth(&mut self, mut n: usize) -> Option<Self::Item> {
    149        if let Some(ref mut a) = self.a {
    150            n = match a.advance_by(n) {
    151                Ok(()) => match a.next() {
    152                    None => 0,
    153                    x => return x,
    154                },
    155                Err(k) => k.get(),
    156            };
    157
    158            self.a = None;
    159        }
    160
    161        self.b.as_mut()?.nth(n)
    162    }
    163
    164    #[inline]
    165    fn find<P>(&mut self, mut predicate: P) -> Option<Self::Item>
    166    where
    167        P: FnMut(&Self::Item) -> bool,
    168    {
    169        and_then_or_clear(&mut self.a, |a| a.find(&mut predicate))
    170            .or_else(|| self.b.as_mut()?.find(predicate))
    171    }
    172
    173    #[inline]
    174    fn last(self) -> Option<A::Item> {
    175        // Must exhaust a before b.
    176        let a_last = self.a.and_then(Iterator::last);
    177        let b_last = self.b.and_then(Iterator::last);
    178        b_last.or(a_last)
    179    }
    180
    181    #[inline]
    182    fn size_hint(&self) -> (usize, Option<usize>) {
    183        match self {
    184            Chain { a: Some(a), b: Some(b) } => {
    185                let (a_lower, a_upper) = a.size_hint();
    186                let (b_lower, b_upper) = b.size_hint();
    187
    188                let lower = a_lower.saturating_add(b_lower);
    189
    190                let upper = match (a_upper, b_upper) {
    191                    (Some(x), Some(y)) => x.checked_add(y),
    192                    _ => None,
    193                };
    194
    195                (lower, upper)
    196            }
    197            Chain { a: Some(a), b: None } => a.size_hint(),
    198            Chain { a: None, b: Some(b) } => b.size_hint(),
    199            Chain { a: None, b: None } => (0, Some(0)),
    200        }
    201    }
    202}
    203
    204#[stable(feature = "rust1", since = "1.0.0")]
    205impl<A, B> DoubleEndedIterator for Chain<A, B>
    206where
    207    A: DoubleEndedIterator,
    208    B: DoubleEndedIterator<Item = A::Item>,
    209{
    210    #[inline]
    211    fn next_back(&mut self) -> Option<A::Item> {
    212        and_then_or_clear(&mut self.b, |b| b.next_back()).or_else(|| self.a.as_mut()?.next_back())
    213    }
    214
    215    #[inline]
    216    fn advance_back_by(&mut self, mut n: usize) -> Result<(), NonZero<usize>> {
    217        if let Some(ref mut b) = self.b {
    218            n = match b.advance_back_by(n) {
    219                Ok(()) => return Ok(()),
    220                Err(k) => k.get(),
    221            };
    222            self.b = None;
    223        }
    224
    225        if let Some(ref mut a) = self.a {
    226            return a.advance_back_by(n);
    227            // we don't fuse the second iterator
    228        }
    229
    230        NonZero::new(n).map_or(Ok(()), Err)
    231    }
    232
    233    #[inline]
    234    fn nth_back(&mut self, mut n: usize) -> Option<Self::Item> {
    235        if let Some(ref mut b) = self.b {
    236            n = match b.advance_back_by(n) {
    237                Ok(()) => match b.next_back() {
    238                    None => 0,
    239                    x => return x,
    240                },
    241                Err(k) => k.get(),
    242            };
    243
    244            self.b = None;
    245        }
    246
    247        self.a.as_mut()?.nth_back(n)
    248    }
    249
    250    #[inline]
    251    fn rfind<P>(&mut self, mut predicate: P) -> Option<Self::Item>
    252    where
    253        P: FnMut(&Self::Item) -> bool,
    254    {
    255        and_then_or_clear(&mut self.b, |b| b.rfind(&mut predicate))
    256            .or_else(|| self.a.as_mut()?.rfind(predicate))
    257    }
    258
    259    fn try_rfold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R
    260    where
    261        Self: Sized,
    262        F: FnMut(Acc, Self::Item) -> R,
    263        R: Try<Output = Acc>,
    264    {
    265        if let Some(ref mut b) = self.b {
    266            acc = b.try_rfold(acc, &mut f)?;
    267            self.b = None;
    268        }
    269        if let Some(ref mut a) = self.a {
    270            acc = a.try_rfold(acc, f)?;
    271            // we don't fuse the second iterator
    272        }
    273        try { acc }
    274    }
    275
    276    fn rfold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc
    277    where
    278        F: FnMut(Acc, Self::Item) -> Acc,
    279    {
    280        if let Some(b) = self.b {
    281            acc = b.rfold(acc, &mut f);
    282        }
    283        if let Some(a) = self.a {
    284            acc = a.rfold(acc, f);
    285        }
    286        acc
    287    }
    288}
    289
    290// Note: *both* must be fused to handle double-ended iterators.
    291#[stable(feature = "fused", since = "1.26.0")]
    292impl<A, B> FusedIterator for Chain<A, B>
    293where
    294    A: FusedIterator,
    295    B: FusedIterator<Item = A::Item>,
    296{
    297}
    298
    299#[unstable(feature = "trusted_len", issue = "37572")]
    300unsafe impl<A, B> TrustedLen for Chain<A, B>
    301where
    302    A: TrustedLen,
    303    B: TrustedLen<Item = A::Item>,
    304{
    305}
    306
    307#[stable(feature = "default_iters", since = "1.70.0")]
    308impl<A: Default, B: Default> Default for Chain<A, B> {
    309    /// Creates a `Chain` from the default values for `A` and `B`.
    310    ///
    311    /// ```
    312    /// # use core::iter::Chain;
    313    /// # use core::slice;
    314    /// # use std::collections::{btree_set, BTreeSet};
    315    /// # use std::mem;
    316    /// struct Foo<'a>(Chain<slice::Iter<'a, u8>, btree_set::Iter<'a, u8>>);
    317    ///
    318    /// let set = BTreeSet::<u8>::new();
    319    /// let slice: &[u8] = &[];
    320    /// let mut foo = Foo(slice.iter().chain(set.iter()));
    321    ///
    322    /// // take requires `Default`
    323    /// let _: Chain<_, _> = mem::take(&mut foo.0);
    324    /// ```
    325    fn default() -> Self {
    326        Chain::new(Default::default(), Default::default())
    327    }
    328}
    329
    330#[inline]
    331fn and_then_or_clear<T, U>(opt: &mut Option<T>, f: impl FnOnce(&mut T) -> Option<U>) -> Option<U> {
    332    let x = f(opt.as_mut()?);
    333    if x.is_none() {
    334        *opt = None;
    335    }
    336    x
    337}