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 ///
  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")]
  23 pub 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 }
  35 impl<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 /// Fuse the iterator if the expression is `None`.
  42 macro_rules! fuse {
  43     ($self:ident . $iter:ident . $($call:tt)+) => {
  44         match $self.$iter {
  45             Some(ref mut iter) => match iter.$($call)+ {
  46                 None => {
  47                     $self.$iter = None;
  48                     None
  49                 }
  50                 item => item,
  51             },
  52             None => None,
  53         }
  54     };
  55 }
  56
  57 /// Try an iterator method without fusing,
  58 /// like an inline `.as_mut().and_then(...)`
  59 macro_rules! maybe {
  60     ($self:ident . $iter:ident . $($call:tt)+) => {
  61         match $self.$iter {
  62             Some(ref mut iter) => iter.$($call)+,
  63             None => None,
  64         }
  65     };
  66 }
  67
  68 #[stable(feature = "rust1", since = "1.0.0")]
  69 impl<A, B> Iterator for Chain<A, B>
  70 where
  71     A: Iterator,
  72     B: Iterator<Item = A::Item>,
  73 {
  74     type Item = A::Item;
  75
  76     #[inline]
  77     fn next(&mut self) -> Option<A::Item> {
  78         match fuse!(self.a.next()) {
  79             None => maybe!(self.b.next()),
  80             item => item,
  81         }
  82     }
  83
  84     #[inline]
  85     #[rustc_inherit_overflow_checks]
  86     fn count(self) -> usize {
  87         let a_count = match self.a {
  88             Some(a) => a.count(),
  89             None => 0,
  90         };
  91         let b_count = match self.b {
  92             Some(b) => b.count(),
  93             None => 0,
  94         };
  95         a_count + b_count
  96     }
  97
  98     fn try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R
  99     where
 100         Self: Sized,
 101         F: FnMut(Acc, Self::Item) -> R,
 102         R: Try<Ok = Acc>,
 103     {
 104         if let Some(ref mut a) = self.a {
 105             acc = a.try_fold(acc, &mut f)?;
 106             self.a = None;
 107         }
 108         if let Some(ref mut b) = self.b {
 109             acc = b.try_fold(acc, f)?;
 110             // we don't fuse the second iterator
 111         }
 112         Try::from_ok(acc)
 113     }
 114
 115     fn fold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc
 116     where
 117         F: FnMut(Acc, Self::Item) -> Acc,
 118     {
 119         if let Some(a) = self.a {
 120             acc = a.fold(acc, &mut f);
 121         }
 122         if let Some(b) = self.b {
 123             acc = b.fold(acc, f);
 124         }
 125         acc
 126     }
 127
 128     #[inline]
 129     fn nth(&mut self, mut n: usize) -> Option<A::Item> {
 130         if let Some(ref mut a) = self.a {
 131             while let Some(x) = a.next() {
 132                 if n == 0 {
 133                     return Some(x);
 134                 }
 135                 n -= 1;
 136             }
 137             self.a = None;
 138         }
 139         maybe!(self.b.nth(n))
 140     }
 141
 142     #[inline]
 143     fn find<P>(&mut self, mut predicate: P) -> Option<Self::Item>
 144     where
 145         P: FnMut(&Self::Item) -> bool,
 146     {
 147         match fuse!(self.a.find(&mut predicate)) {
 148             None => maybe!(self.b.find(predicate)),
 149             item => item,
 150         }
 151     }
 152
 153     #[inline]
 154     fn last(self) -> Option<A::Item> {
 155         // Must exhaust a before b.
 156         let a_last = match self.a {
 157             Some(a) => a.last(),
 158             None => None,
 159         };
 160         let b_last = match self.b {
 161             Some(b) => b.last(),
 162             None => None,
 163         };
 164         b_last.or(a_last)
 165     }
 166
 167     #[inline]
 168     fn size_hint(&self) -> (usize, Option<usize>) {
 169         match self {
 170             Chain { a: Some(a), b: Some(b) } => {
 171                 let (a_lower, a_upper) = a.size_hint();
 172                 let (b_lower, b_upper) = b.size_hint();
 173
 174                 let lower = a_lower.saturating_add(b_lower);
 175
 176                 let upper = match (a_upper, b_upper) {
 177                     (Some(x), Some(y)) => x.checked_add(y),
 178                     _ => None,
 179                 };
 180
 181                 (lower, upper)
 182             }
 183             Chain { a: Some(a), b: None } => a.size_hint(),
 184             Chain { a: None, b: Some(b) } => b.size_hint(),
 185             Chain { a: None, b: None } => (0, Some(0)),
 186         }
 187     }
 188 }
 189
 190 #[stable(feature = "rust1", since = "1.0.0")]
 191 impl<A, B> DoubleEndedIterator for Chain<A, B>
 192 where
 193     A: DoubleEndedIterator,
 194     B: DoubleEndedIterator<Item = A::Item>,
 195 {
 196     #[inline]
 197     fn next_back(&mut self) -> Option<A::Item> {
 198         match fuse!(self.b.next_back()) {
 199             None => maybe!(self.a.next_back()),
 200             item => item,
 201         }
 202     }
 203
 204     #[inline]
 205     fn nth_back(&mut self, mut n: usize) -> Option<A::Item> {
 206         if let Some(ref mut b) = self.b {
 207             while let Some(x) = b.next_back() {
 208                 if n == 0 {
 209                     return Some(x);
 210                 }
 211                 n -= 1;
 212             }
 213             self.b = None;
 214         }
 215         maybe!(self.a.nth_back(n))
 216     }
 217
 218     #[inline]
 219     fn rfind<P>(&mut self, mut predicate: P) -> Option<Self::Item>
 220     where
 221         P: FnMut(&Self::Item) -> bool,
 222     {
 223         match fuse!(self.b.rfind(&mut predicate)) {
 224             None => maybe!(self.a.rfind(predicate)),
 225             item => item,
 226         }
 227     }
 228
 229     fn try_rfold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R
 230     where
 231         Self: Sized,
 232         F: FnMut(Acc, Self::Item) -> R,
 233         R: Try<Ok = Acc>,
 234     {
 235         if let Some(ref mut b) = self.b {
 236             acc = b.try_rfold(acc, &mut f)?;
 237             self.b = None;
 238         }
 239         if let Some(ref mut a) = self.a {
 240             acc = a.try_rfold(acc, f)?;
 241             // we don't fuse the second iterator
 242         }
 243         Try::from_ok(acc)
 244     }
 245
 246     fn rfold<Acc, F>(self, mut acc: Acc, mut f: F) -> Acc
 247     where
 248         F: FnMut(Acc, Self::Item) -> Acc,
 249     {
 250         if let Some(b) = self.b {
 251             acc = b.rfold(acc, &mut f);
 252         }
 253         if let Some(a) = self.a {
 254             acc = a.rfold(acc, f);
 255         }
 256         acc
 257     }
 258 }
 259
 260 // Note: *both* must be fused to handle double-ended iterators.
 261 #[stable(feature = "fused", since = "1.26.0")]
 262 impl<A, B> FusedIterator for Chain<A, B>
 263 where
 264     A: FusedIterator,
 265     B: FusedIterator<Item = A::Item>,
 266 {
 267 }
 268
 269 #[unstable(feature = "trusted_len", issue = "37572")]
 270 unsafe impl<A, B> TrustedLen for Chain<A, B>
 271 where
 272     A: TrustedLen,
 273     B: TrustedLen<Item = A::Item>,
 274 {
 275 }