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

   1 use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedLen};
   2 use crate::ops::Try;
   3
   4 /// An iterator with a `peek()` that returns an optional reference to the next
   5 /// element.
   6 ///
   7 /// This `struct` is created by the [`peekable`] method on [`Iterator`]. See its
   8 /// documentation for more.
   9 ///
  10 /// [`peekable`]: Iterator::peekable
  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 Peekable<I: Iterator> {
  16     iter: I,
  17     /// Remember a peeked value, even if it was None.
  18     peeked: Option<Option<I::Item>>,
  19 }
  20
  21 impl<I: Iterator> Peekable<I> {
  22     pub(in crate::iter) fn new(iter: I) -> Peekable<I> {
  23         Peekable { iter, peeked: None }
  24     }
  25 }
  26
  27 // Peekable must remember if a None has been seen in the `.peek()` method.
  28 // It ensures that `.peek(); .peek();` or `.peek(); .next();` only advances the
  29 // underlying iterator at most once. This does not by itself make the iterator
  30 // fused.
  31 #[stable(feature = "rust1", since = "1.0.0")]
  32 impl<I: Iterator> Iterator for Peekable<I> {
  33     type Item = I::Item;
  34
  35     #[inline]
  36     fn next(&mut self) -> Option<I::Item> {
  37         match self.peeked.take() {
  38             Some(v) => v,
  39             None => self.iter.next(),
  40         }
  41     }
  42
  43     #[inline]
  44     #[rustc_inherit_overflow_checks]
  45     fn count(mut self) -> usize {
  46         match self.peeked.take() {
  47             Some(None) => 0,
  48             Some(Some(_)) => 1 + self.iter.count(),
  49             None => self.iter.count(),
  50         }
  51     }
  52
  53     #[inline]
  54     fn nth(&mut self, n: usize) -> Option<I::Item> {
  55         match self.peeked.take() {
  56             Some(None) => None,
  57             Some(v @ Some(_)) if n == 0 => v,
  58             Some(Some(_)) => self.iter.nth(n - 1),
  59             None => self.iter.nth(n),
  60         }
  61     }
  62
  63     #[inline]
  64     fn last(mut self) -> Option<I::Item> {
  65         let peek_opt = match self.peeked.take() {
  66             Some(None) => return None,
  67             Some(v) => v,
  68             None => None,
  69         };
  70         self.iter.last().or(peek_opt)
  71     }
  72
  73     #[inline]
  74     fn size_hint(&self) -> (usize, Option<usize>) {
  75         let peek_len = match self.peeked {
  76             Some(None) => return (0, Some(0)),
  77             Some(Some(_)) => 1,
  78             None => 0,
  79         };
  80         let (lo, hi) = self.iter.size_hint();
  81         let lo = lo.saturating_add(peek_len);
  82         let hi = match hi {
  83             Some(x) => x.checked_add(peek_len),
  84             None => None,
  85         };
  86         (lo, hi)
  87     }
  88
  89     #[inline]
  90     fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R
  91     where
  92         Self: Sized,
  93         F: FnMut(B, Self::Item) -> R,
  94         R: Try<Ok = B>,
  95     {
  96         let acc = match self.peeked.take() {
  97             Some(None) => return try { init },
  98             Some(Some(v)) => f(init, v)?,
  99             None => init,
 100         };
 101         self.iter.try_fold(acc, f)
 102     }
 103
 104     #[inline]
 105     fn fold<Acc, Fold>(self, init: Acc, mut fold: Fold) -> Acc
 106     where
 107         Fold: FnMut(Acc, Self::Item) -> Acc,
 108     {
 109         let acc = match self.peeked {
 110             Some(None) => return init,
 111             Some(Some(v)) => fold(init, v),
 112             None => init,
 113         };
 114         self.iter.fold(acc, fold)
 115     }
 116 }
 117
 118 #[stable(feature = "double_ended_peek_iterator", since = "1.38.0")]
 119 impl<I> DoubleEndedIterator for Peekable<I>
 120 where
 121     I: DoubleEndedIterator,
 122 {
 123     #[inline]
 124     fn next_back(&mut self) -> Option<Self::Item> {
 125         match self.peeked.as_mut() {
 126             Some(v @ Some(_)) => self.iter.next_back().or_else(|| v.take()),
 127             Some(None) => None,
 128             None => self.iter.next_back(),
 129         }
 130     }
 131
 132     #[inline]
 133     fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R
 134     where
 135         Self: Sized,
 136         F: FnMut(B, Self::Item) -> R,
 137         R: Try<Ok = B>,
 138     {
 139         match self.peeked.take() {
 140             Some(None) => try { init },
 141             Some(Some(v)) => match self.iter.try_rfold(init, &mut f).into_result() {
 142                 Ok(acc) => f(acc, v),
 143                 Err(e) => {
 144                     self.peeked = Some(Some(v));
 145                     Try::from_error(e)
 146                 }
 147             },
 148             None => self.iter.try_rfold(init, f),
 149         }
 150     }
 151
 152     #[inline]
 153     fn rfold<Acc, Fold>(self, init: Acc, mut fold: Fold) -> Acc
 154     where
 155         Fold: FnMut(Acc, Self::Item) -> Acc,
 156     {
 157         match self.peeked {
 158             Some(None) => init,
 159             Some(Some(v)) => {
 160                 let acc = self.iter.rfold(init, &mut fold);
 161                 fold(acc, v)
 162             }
 163             None => self.iter.rfold(init, fold),
 164         }
 165     }
 166 }
 167
 168 #[stable(feature = "rust1", since = "1.0.0")]
 169 impl<I: ExactSizeIterator> ExactSizeIterator for Peekable<I> {}
 170
 171 #[stable(feature = "fused", since = "1.26.0")]
 172 impl<I: FusedIterator> FusedIterator for Peekable<I> {}
 173
 174 impl<I: Iterator> Peekable<I> {
 175     /// Returns a reference to the next() value without advancing the iterator.
 176     ///
 177     /// Like [`next`], if there is a value, it is wrapped in a `Some(T)`.
 178     /// But if the iteration is over, `None` is returned.
 179     ///
 180     /// [`next`]: Iterator::next
 181     ///
 182     /// Because `peek()` returns a reference, and many iterators iterate over
 183     /// references, there can be a possibly confusing situation where the
 184     /// return value is a double reference. You can see this effect in the
 185     /// examples below.
 186     ///
 187     /// # Examples
 188     ///
 189     /// Basic usage:
 190     ///
 191     /// ```
 192     /// let xs = [1, 2, 3];
 193     ///
 194     /// let mut iter = xs.iter().peekable();
 195     ///
 196     /// // peek() lets us see into the future
 197     /// assert_eq!(iter.peek(), Some(&&1));
 198     /// assert_eq!(iter.next(), Some(&1));
 199     ///
 200     /// assert_eq!(iter.next(), Some(&2));
 201     ///
 202     /// // The iterator does not advance even if we `peek` multiple times
 203     /// assert_eq!(iter.peek(), Some(&&3));
 204     /// assert_eq!(iter.peek(), Some(&&3));
 205     ///
 206     /// assert_eq!(iter.next(), Some(&3));
 207     ///
 208     /// // After the iterator is finished, so is `peek()`
 209     /// assert_eq!(iter.peek(), None);
 210     /// assert_eq!(iter.next(), None);
 211     /// ```
 212     #[inline]
 213     #[stable(feature = "rust1", since = "1.0.0")]
 214     pub fn peek(&mut self) -> Option<&I::Item> {
 215         let iter = &mut self.iter;
 216         self.peeked.get_or_insert_with(|| iter.next()).as_ref()
 217     }
 218
 219     /// Returns a mutable reference to the next() value without advancing the iterator.
 220     ///
 221     /// Like [`next`], if there is a value, it is wrapped in a `Some(T)`.
 222     /// But if the iteration is over, `None` is returned.
 223     ///
 224     /// Because `peek_mut()` returns a reference, and many iterators iterate over
 225     /// references, there can be a possibly confusing situation where the
 226     /// return value is a double reference. You can see this effect in the examples
 227     /// below.
 228     ///
 229     /// [`next`]: Iterator::next
 230     ///
 231     /// # Examples
 232     ///
 233     /// ```
 234     /// #![feature(peekable_peek_mut)]
 235     /// let mut iter = [1, 2, 3].iter().peekable();
 236     ///
 237     /// assert_eq!(iter.peek_mut(), Some(&mut &1));
 238     /// assert_eq!(iter.next(), Some(&1));
 239     ///
 240     /// // Peek into the iterator and modify the value which will be returned next
 241     /// if let Some(mut p) = iter.peek_mut() {
 242     ///     if *p == &2 {
 243     ///         *p = &5;
 244     ///     }
 245     /// }
 246     ///
 247     /// assert_eq!(iter.collect::<Vec<_>>(), vec![&5, &3]);
 248     /// ```
 249     #[inline]
 250     #[unstable(feature = "peekable_peek_mut", issue = "78302")]
 251     pub fn peek_mut(&mut self) -> Option<&mut I::Item> {
 252         let iter = &mut self.iter;
 253         self.peeked.get_or_insert_with(|| iter.next()).as_mut()
 254     }
 255
 256     /// Consume and return the next value of this iterator if a condition is true.
 257     ///
 258     /// If `func` returns `true` for the next value of this iterator, consume and return it.
 259     /// Otherwise, return `None`.
 260     ///
 261     /// # Examples
 262     /// Consume a number if it's equal to 0.
 263     /// ```
 264     /// #![feature(peekable_next_if)]
 265     /// let mut iter = (0..5).peekable();
 266     /// // The first item of the iterator is 0; consume it.
 267     /// assert_eq!(iter.next_if(|&x| x == 0), Some(0));
 268     /// // The next item returned is now 1, so `consume` will return `false`.
 269     /// assert_eq!(iter.next_if(|&x| x == 0), None);
 270     /// // `next_if` saves the value of the next item if it was not equal to `expected`.
 271     /// assert_eq!(iter.next(), Some(1));
 272     /// ```
 273     ///
 274     /// Consume any number less than 10.
 275     /// ```
 276     /// #![feature(peekable_next_if)]
 277     /// let mut iter = (1..20).peekable();
 278     /// // Consume all numbers less than 10
 279     /// while iter.next_if(|&x| x < 10).is_some() {}
 280     /// // The next value returned will be 10
 281     /// assert_eq!(iter.next(), Some(10));
 282     /// ```
 283     #[unstable(feature = "peekable_next_if", issue = "72480")]
 284     pub fn next_if(&mut self, func: impl FnOnce(&I::Item) -> bool) -> Option<I::Item> {
 285         match self.next() {
 286             Some(matched) if func(&matched) => Some(matched),
 287             other => {
 288                 // Since we called `self.next()`, we consumed `self.peeked`.
 289                 assert!(self.peeked.is_none());
 290                 self.peeked = Some(other);
 291                 None
 292             }
 293         }
 294     }
 295
 296     /// Consume and return the next item if it is equal to `expected`.
 297     ///
 298     /// # Example
 299     /// Consume a number if it's equal to 0.
 300     /// ```
 301     /// #![feature(peekable_next_if)]
 302     /// let mut iter = (0..5).peekable();
 303     /// // The first item of the iterator is 0; consume it.
 304     /// assert_eq!(iter.next_if_eq(&0), Some(0));
 305     /// // The next item returned is now 1, so `consume` will return `false`.
 306     /// assert_eq!(iter.next_if_eq(&0), None);
 307     /// // `next_if_eq` saves the value of the next item if it was not equal to `expected`.
 308     /// assert_eq!(iter.next(), Some(1));
 309     /// ```
 310     #[unstable(feature = "peekable_next_if", issue = "72480")]
 311     pub fn next_if_eq<T>(&mut self, expected: &T) -> Option<I::Item>
 312     where
 313         T: ?Sized,
 314         I::Item: PartialEq<T>,
 315     {
 316         self.next_if(|next| next == expected)
 317     }
 318 }
 319
 320 #[unstable(feature = "trusted_len", issue = "37572")]
 321 unsafe impl<I> TrustedLen for Peekable<I> where I: TrustedLen {}
 322
 323 #[unstable(issue = "none", feature = "inplace_iteration")]
 324 unsafe impl<S: Iterator, I: Iterator> SourceIter for Peekable<I>
 325 where
 326     I: SourceIter<Source = S>,
 327 {
 328     type Source = S;
 329
 330     #[inline]
 331     unsafe fn as_inner(&mut self) -> &mut S {
 332         // SAFETY: unsafe function forwarding to unsafe function with the same requirements
 333         unsafe { SourceIter::as_inner(&mut self.iter) }
 334     }
 335 }
 336
 337 #[unstable(issue = "none", feature = "inplace_iteration")]
 338 unsafe impl<I: InPlaceIterable> InPlaceIterable for Peekable<I> {}