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

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