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