src/libcore/iter/adapters/chain.rs

   1 use crate::ops::Try;
   2 use crate::usize;
   3
   4 use super::super::{Iterator, DoubleEndedIterator, FusedIterator, TrustedLen};
   5
   6 /// An iterator that links two iterators together, in a chain.
   7 ///
   8 /// This `struct` is created by the [`chain`] method on [`Iterator`]. See its
   9 /// documentation for more.
  10 ///
  11 /// [`chain`]: trait.Iterator.html#method.chain
  12 /// [`Iterator`]: trait.Iterator.html
  13 #[derive(Clone, Debug)]
  14 #[must_use = "iterators are lazy and do nothing unless consumed"]
  15 #[stable(feature = "rust1", since = "1.0.0")]
  16 pub struct Chain<A, B> {
  17     a: A,
  18     b: B,
  19     state: ChainState,
  20 }
  21 impl<A, B> Chain<A, B> {
  22     pub(in super::super) fn new(a: A, b: B) -> Chain<A, B> {
  23         Chain { a, b, state: ChainState::Both }
  24     }
  25 }
  26
  27 // The iterator protocol specifies that iteration ends with the return value
  28 // `None` from `.next()` (or `.next_back()`) and it is unspecified what
  29 // further calls return. The chain adaptor must account for this since it uses
  30 // two subiterators.
  31 //
  32 //  It uses three states:
  33 //
  34 //  - Both: `a` and `b` are remaining
  35 //  - Front: `a` remaining
  36 //  - Back: `b` remaining
  37 //
  38 //  The fourth state (neither iterator is remaining) only occurs after Chain has
  39 //  returned None once, so we don't need to store this state.
  40 #[derive(Clone, Debug)]
  41 enum ChainState {
  42     // both front and back iterator are remaining
  43     Both,
  44     // only front is remaining
  45     Front,
  46     // only back is remaining
  47     Back,
  48 }
  49
  50 #[stable(feature = "rust1", since = "1.0.0")]
  51 impl<A, B> Iterator for Chain<A, B> where
  52     A: Iterator,
  53     B: Iterator<Item = A::Item>
  54 {
  55     type Item = A::Item;
  56
  57     #[inline]
  58     fn next(&mut self) -> Option<A::Item> {
  59         match self.state {
  60             ChainState::Both => match self.a.next() {
  61                 elt @ Some(..) => elt,
  62                 None => {
  63                     self.state = ChainState::Back;
  64                     self.b.next()
  65                 }
  66             },
  67             ChainState::Front => self.a.next(),
  68             ChainState::Back => self.b.next(),
  69         }
  70     }
  71
  72     #[inline]
  73     #[rustc_inherit_overflow_checks]
  74     fn count(self) -> usize {
  75         match self.state {
  76             ChainState::Both => self.a.count() + self.b.count(),
  77             ChainState::Front => self.a.count(),
  78             ChainState::Back => self.b.count(),
  79         }
  80     }
  81
  82     fn try_fold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R where
  83         Self: Sized, F: FnMut(Acc, Self::Item) -> R, R: Try<Ok=Acc>
  84     {
  85         let mut accum = init;
  86         match self.state {
  87             ChainState::Both | ChainState::Front => {
  88                 accum = self.a.try_fold(accum, &mut f)?;
  89                 if let ChainState::Both = self.state {
  90                     self.state = ChainState::Back;
  91                 }
  92             }
  93             _ => { }
  94         }
  95         if let ChainState::Back = self.state {
  96             accum = self.b.try_fold(accum, &mut f)?;
  97         }
  98         Try::from_ok(accum)
  99     }
 100
 101     fn fold<Acc, F>(self, init: Acc, mut f: F) -> Acc
 102         where F: FnMut(Acc, Self::Item) -> Acc,
 103     {
 104         let mut accum = init;
 105         match self.state {
 106             ChainState::Both | ChainState::Front => {
 107                 accum = self.a.fold(accum, &mut f);
 108             }
 109             _ => { }
 110         }
 111         match self.state {
 112             ChainState::Both | ChainState::Back => {
 113                 accum = self.b.fold(accum, &mut f);
 114             }
 115             _ => { }
 116         }
 117         accum
 118     }
 119
 120     #[inline]
 121     fn nth(&mut self, mut n: usize) -> Option<A::Item> {
 122         match self.state {
 123             ChainState::Both | ChainState::Front => {
 124                 for x in self.a.by_ref() {
 125                     if n == 0 {
 126                         return Some(x)
 127                     }
 128                     n -= 1;
 129                 }
 130                 if let ChainState::Both = self.state {
 131                     self.state = ChainState::Back;
 132                 }
 133             }
 134             ChainState::Back => {}
 135         }
 136         if let ChainState::Back = self.state {
 137             self.b.nth(n)
 138         } else {
 139             None
 140         }
 141     }
 142
 143     #[inline]
 144     fn find<P>(&mut self, mut predicate: P) -> Option<Self::Item> where
 145         P: FnMut(&Self::Item) -> bool,
 146     {
 147         match self.state {
 148             ChainState::Both => match self.a.find(&mut predicate) {
 149                 None => {
 150                     self.state = ChainState::Back;
 151                     self.b.find(predicate)
 152                 }
 153                 v => v
 154             },
 155             ChainState::Front => self.a.find(predicate),
 156             ChainState::Back => self.b.find(predicate),
 157         }
 158     }
 159
 160     #[inline]
 161     fn last(self) -> Option<A::Item> {
 162         match self.state {
 163             ChainState::Both => {
 164                 // Must exhaust a before b.
 165                 let a_last = self.a.last();
 166                 let b_last = self.b.last();
 167                 b_last.or(a_last)
 168             },
 169             ChainState::Front => self.a.last(),
 170             ChainState::Back => self.b.last()
 171         }
 172     }
 173
 174     #[inline]
 175     fn size_hint(&self) -> (usize, Option<usize>) {
 176         match self.state {
 177             ChainState::Both => {
 178                 let (a_lower, a_upper) = self.a.size_hint();
 179                 let (b_lower, b_upper) = self.b.size_hint();
 180
 181                 let lower = a_lower.saturating_add(b_lower);
 182
 183                 let upper = match (a_upper, b_upper) {
 184                     (Some(x), Some(y)) => x.checked_add(y),
 185                     _ => None
 186                 };
 187
 188                 (lower, upper)
 189             }
 190             ChainState::Front => self.a.size_hint(),
 191             ChainState::Back => self.b.size_hint(),
 192         }
 193     }
 194 }
 195
 196 #[stable(feature = "rust1", since = "1.0.0")]
 197 impl<A, B> DoubleEndedIterator for Chain<A, B> where
 198     A: DoubleEndedIterator,
 199     B: DoubleEndedIterator<Item=A::Item>,
 200 {
 201     #[inline]
 202     fn next_back(&mut self) -> Option<A::Item> {
 203         match self.state {
 204             ChainState::Both => match self.b.next_back() {
 205                 elt @ Some(..) => elt,
 206                 None => {
 207                     self.state = ChainState::Front;
 208                     self.a.next_back()
 209                 }
 210             },
 211             ChainState::Front => self.a.next_back(),
 212             ChainState::Back => self.b.next_back(),
 213         }
 214     }
 215
 216     #[inline]
 217     fn nth_back(&mut self, mut n: usize) -> Option<A::Item> {
 218         match self.state {
 219             ChainState::Both | ChainState::Back => {
 220                 for x in self.b.by_ref().rev() {
 221                     if n == 0 {
 222                         return Some(x)
 223                     }
 224                     n -= 1;
 225                 }
 226                 if let ChainState::Both = self.state {
 227                     self.state = ChainState::Front;
 228                 }
 229             }
 230             ChainState::Front => {}
 231         }
 232         if let ChainState::Front = self.state {
 233             self.a.nth_back(n)
 234         } else {
 235             None
 236         }
 237     }
 238
 239     fn try_rfold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R where
 240         Self: Sized, F: FnMut(Acc, Self::Item) -> R, R: Try<Ok=Acc>
 241     {
 242         let mut accum = init;
 243         match self.state {
 244             ChainState::Both | ChainState::Back => {
 245                 accum = self.b.try_rfold(accum, &mut f)?;
 246                 if let ChainState::Both = self.state {
 247                     self.state = ChainState::Front;
 248                 }
 249             }
 250             _ => { }
 251         }
 252         if let ChainState::Front = self.state {
 253             accum = self.a.try_rfold(accum, &mut f)?;
 254         }
 255         Try::from_ok(accum)
 256     }
 257
 258     fn rfold<Acc, F>(self, init: Acc, mut f: F) -> Acc
 259         where F: FnMut(Acc, Self::Item) -> Acc,
 260     {
 261         let mut accum = init;
 262         match self.state {
 263             ChainState::Both | ChainState::Back => {
 264                 accum = self.b.rfold(accum, &mut f);
 265             }
 266             _ => { }
 267         }
 268         match self.state {
 269             ChainState::Both | ChainState::Front => {
 270                 accum = self.a.rfold(accum, &mut f);
 271             }
 272             _ => { }
 273         }
 274         accum
 275     }
 276
 277 }
 278
 279 // Note: *both* must be fused to handle double-ended iterators.
 280 #[stable(feature = "fused", since = "1.26.0")]
 281 impl<A, B> FusedIterator for Chain<A, B>
 282     where A: FusedIterator,
 283           B: FusedIterator<Item=A::Item>,
 284 {}
 285
 286 #[unstable(feature = "trusted_len", issue = "37572")]
 287 unsafe impl<A, B> TrustedLen for Chain<A, B>
 288     where A: TrustedLen, B: TrustedLen<Item=A::Item>,
 289 {}