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` retains the next item if the predicate evaluates to `false` for it.
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_eq` will return `None`.
308 /// assert_eq!(iter.next_if_eq(&0), None);
309 /// // `next_if_eq` retains 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}
338