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