core/iter/adapters/
peekable.rs

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