library/core/src/iter/adapters/chain.rs

   1 use crate::iter::{DoubleEndedIterator, FusedIterator, Iterator, TrustedLen};
   2 use crate::ops::Try;
   3 use crate::usize;
   4
   5 /// An iterator that links two iterators together, in a chain.
   6 ///
   7 /// This `struct` is created by [`Iterator::chain`]. See its documentation
   8 /// for more.
   9 #[derive(Clone, Debug)]
  10 #[must_use = "iterators are lazy and do nothing unless consumed"]
  11 #[stable(feature = "rust1", since = "1.0.0")]
  12 pub struct Chain<A, B> {
  13     // These are "fused" with `Option` so we don't need separate state to track which part is
  14     // already exhausted, and we may also get niche layout for `None`. We don't use the real `Fuse`
  15     // adapter because its specialization for `FusedIterator` unconditionally descends into the
  16     // iterator, and that could be expensive to keep revisiting stuff like nested chains. It also
  17     // hurts compiler performance to add more iterator layers to `Chain`.
  18     //
  19     // Only the "first" iterator is actually set `None` when exhausted, depending on whether you
  20     // iterate forward or backward. If you mix directions, then both sides may be `None`.
  21     a: Option<A>,
  22     b: Option<B>,
  23 }
  24 impl<A, B> Chain<A, B> {
  25     pub(in super::super) fn new(a: A, b: B) -> Chain<A, B> {
  26         Chain { a: Some(a), b: Some(b) }
  27     }
  28 }
  29
  30 /// Fuse the iterator if the expression is `None`.
  31 macro_rules! fuse {
  32     ($self:ident . $iter:ident . $($call:tt)+) => {
  33         match $self.$iter {
  34             Some(ref mut iter) => match iter.$($call)+ {
  35                 None => {
  36                     $self.$iter = None;
  37                     None
  38                 }
  39                 item => item,
  40             },
  41             None => None,
  42         }
  43     };
  44 }
  45
  46 /// Try an iterator method without fusing,
  47 /// like an inline `.as_mut().and_then(...)`
  48 macro_rules! maybe {
  49     ($self:ident . $iter:ident . $($call:tt)+) => {
  50         match $self.$iter {
  51             Some(ref mut iter) => iter.$($call)+,
  52             None => None,
  53         }
  54     };
  55 }
  56
  57 #[stable(feature = "rust1", since = "1.0.0")]
  58 impl<A, B> Iterator for Chain<A, B>
  59 where
  60     A: Iterator,
  61     B: Iterator<Item = A::Item>,
  62 {
  63     type Item = A::Item;
  64
  65     #[inline]
  66     fn next(&mut self) -> Option<A::Item> {
  67         match fuse!(self.a.next()) {
  68             None => maybe!(self.b.next()),
  69             item => item,
  70         }
  71     }
  72
  73     #[inline]
  74     #[rustc_inherit_overflow_checks]
  75     fn count(self) -> usize {
  76         let a_count = match self.a {
  77             Some(a) => a.count(),
  78             None => 0,
  79         };
  80         let b_count = match self.b {
  81             Some(b) => b.count(),
  82             None => 0,
  83         };
  84         a_count + b_count
  85     }
  86
  87     fn try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R
  88     where
  89         Self: Sized,
  90         F: FnMut(Acc, Self::Item) -> R,
  91         R: Try<Ok = Acc>,
  92     {
  93         if let Some(ref mut a) = self.a {
  94             acc = a.try_fold(acc, &mut f)?;
  95             self.a = None;
  96         }
  97         if let Some(ref mut b) = self.b {
  98             acc = b.try_fold(acc, f)?;
  99             // we don't fuse the second iterator
 100         }
 101         Try::from_ok(acc)
 102     }
 103
 104     fn fold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc
 105     where
 106         F: FnMut(Acc, Self::Item) -> Acc,
 107     {
 108         if let Some(a) = self.a {
 109             acc = a.fold(acc, &mut f);
 110         }
 111         if let Some(b) = self.b {
 112             acc = b.fold(acc, f);
 113         }
 114         acc
 115     }
 116
 117     #[inline]
 118     fn nth(&mut self, mut n: usize) -> Option<A::Item> {
 119         if let Some(ref mut a) = self.a {
 120             while let Some(x) = a.next() {
 121                 if n == 0 {
 122                     return Some(x);
 123                 }
 124                 n -= 1;
 125             }
 126             self.a = None;
 127         }
 128         maybe!(self.b.nth(n))
 129     }
 130
 131     #[inline]
 132     fn find<P>(&mut self, mut predicate: P) -> Option<Self::Item>
 133     where
 134         P: FnMut(&Self::Item) -> bool,
 135     {
 136         match fuse!(self.a.find(&mut predicate)) {
 137             None => maybe!(self.b.find(predicate)),
 138             item => item,
 139         }
 140     }
 141
 142     #[inline]
 143     fn last(self) -> Option<A::Item> {
 144         // Must exhaust a before b.
 145         let a_last = match self.a {
 146             Some(a) => a.last(),
 147             None => None,
 148         };
 149         let b_last = match self.b {
 150             Some(b) => b.last(),
 151             None => None,
 152         };
 153         b_last.or(a_last)
 154     }
 155
 156     #[inline]
 157     fn size_hint(&self) -> (usize, Option<usize>) {
 158         match self {
 159             Chain { a: Some(a), b: Some(b) } => {
 160                 let (a_lower, a_upper) = a.size_hint();
 161                 let (b_lower, b_upper) = b.size_hint();
 162
 163                 let lower = a_lower.saturating_add(b_lower);
 164
 165                 let upper = match (a_upper, b_upper) {
 166                     (Some(x), Some(y)) => x.checked_add(y),
 167                     _ => None,
 168                 };
 169
 170                 (lower, upper)
 171             }
 172             Chain { a: Some(a), b: None } => a.size_hint(),
 173             Chain { a: None, b: Some(b) } => b.size_hint(),
 174             Chain { a: None, b: None } => (0, Some(0)),
 175         }
 176     }
 177 }
 178
 179 #[stable(feature = "rust1", since = "1.0.0")]
 180 impl<A, B> DoubleEndedIterator for Chain<A, B>
 181 where
 182     A: DoubleEndedIterator,
 183     B: DoubleEndedIterator<Item = A::Item>,
 184 {
 185     #[inline]
 186     fn next_back(&mut self) -> Option<A::Item> {
 187         match fuse!(self.b.next_back()) {
 188             None => maybe!(self.a.next_back()),
 189             item => item,
 190         }
 191     }
 192
 193     #[inline]
 194     fn nth_back(&mut self, mut n: usize) -> Option<A::Item> {
 195         if let Some(ref mut b) = self.b {
 196             while let Some(x) = b.next_back() {
 197                 if n == 0 {
 198                     return Some(x);
 199                 }
 200                 n -= 1;
 201             }
 202             self.b = None;
 203         }
 204         maybe!(self.a.nth_back(n))
 205     }
 206
 207     #[inline]
 208     fn rfind<P>(&mut self, mut predicate: P) -> Option<Self::Item>
 209     where
 210         P: FnMut(&Self::Item) -> bool,
 211     {
 212         match fuse!(self.b.rfind(&mut predicate)) {
 213             None => maybe!(self.a.rfind(predicate)),
 214             item => item,
 215         }
 216     }
 217
 218     fn try_rfold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R
 219     where
 220         Self: Sized,
 221         F: FnMut(Acc, Self::Item) -> R,
 222         R: Try<Ok = Acc>,
 223     {
 224         if let Some(ref mut b) = self.b {
 225             acc = b.try_rfold(acc, &mut f)?;
 226             self.b = None;
 227         }
 228         if let Some(ref mut a) = self.a {
 229             acc = a.try_rfold(acc, f)?;
 230             // we don't fuse the second iterator
 231         }
 232         Try::from_ok(acc)
 233     }
 234
 235     fn rfold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc
 236     where
 237         F: FnMut(Acc, Self::Item) -> Acc,
 238     {
 239         if let Some(b) = self.b {
 240             acc = b.rfold(acc, &mut f);
 241         }
 242         if let Some(a) = self.a {
 243             acc = a.rfold(acc, f);
 244         }
 245         acc
 246     }
 247 }
 248
 249 // Note: *both* must be fused to handle double-ended iterators.
 250 #[stable(feature = "fused", since = "1.26.0")]
 251 impl<A, B> FusedIterator for Chain<A, B>
 252 where
 253     A: FusedIterator,
 254     B: FusedIterator<Item = A::Item>,
 255 {
 256 }
 257
 258 #[unstable(feature = "trusted_len", issue = "37572")]
 259 unsafe impl<A, B> TrustedLen for Chain<A, B>
 260 where
 261     A: TrustedLen,
 262     B: TrustedLen<Item = A::Item>,
 263 {
 264 }