]>
Commit | Line | Data |
---|---|---|
cdc7bbd5 | 1 | use crate::iter::{adapters::SourceIter, FusedIterator, TrustedLen}; |
17df50a5 | 2 | use crate::ops::{ControlFlow, Try}; |
fc512014 XL |
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, | |
17df50a5 | 94 | R: Try<Output = B>, |
fc512014 XL |
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, | |
17df50a5 | 137 | R: Try<Output = B>, |
fc512014 | 138 | { |
17df50a5 XL |
139 | match self.peeked.take() { |
140 | Some(None) => try { init }, | |
141 | Some(Some(v)) => match self.iter.try_rfold(init, &mut f).branch() { | |
142 | ControlFlow::Continue(acc) => f(acc, v), | |
143 | ControlFlow::Break(r) => { | |
144 | self.peeked = Some(Some(v)); | |
145 | R::from_residual(r) | |
146 | } | |
147 | }, | |
148 | None => self.iter.try_rfold(init, f), | |
149 | } | |
150 | } | |
151 | ||
fc512014 XL |
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 | /// Basic usage: | |
234 | /// | |
235 | /// ``` | |
fc512014 XL |
236 | /// let mut iter = [1, 2, 3].iter().peekable(); |
237 | /// | |
238 | /// // Like with `peek()`, we can see into the future without advancing the iterator. | |
239 | /// assert_eq!(iter.peek_mut(), Some(&mut &1)); | |
240 | /// assert_eq!(iter.peek_mut(), Some(&mut &1)); | |
241 | /// assert_eq!(iter.next(), Some(&1)); | |
242 | /// | |
243 | /// // Peek into the iterator and set the value behind the mutable reference. | |
244 | /// if let Some(p) = iter.peek_mut() { | |
245 | /// assert_eq!(*p, &2); | |
246 | /// *p = &5; | |
247 | /// } | |
248 | /// | |
249 | /// // The value we put in reappears as the iterator continues. | |
250 | /// assert_eq!(iter.collect::<Vec<_>>(), vec![&5, &3]); | |
251 | /// ``` | |
252 | #[inline] | |
cdc7bbd5 | 253 | #[stable(feature = "peekable_peek_mut", since = "1.53.0")] |
fc512014 XL |
254 | pub fn peek_mut(&mut self) -> Option<&mut I::Item> { |
255 | let iter = &mut self.iter; | |
256 | self.peeked.get_or_insert_with(|| iter.next()).as_mut() | |
257 | } | |
258 | ||
259 | /// Consume and return the next value of this iterator if a condition is true. | |
260 | /// | |
261 | /// If `func` returns `true` for the next value of this iterator, consume and return it. | |
262 | /// Otherwise, return `None`. | |
263 | /// | |
264 | /// # Examples | |
265 | /// Consume a number if it's equal to 0. | |
266 | /// ``` | |
fc512014 XL |
267 | /// let mut iter = (0..5).peekable(); |
268 | /// // The first item of the iterator is 0; consume it. | |
269 | /// assert_eq!(iter.next_if(|&x| x == 0), Some(0)); | |
270 | /// // The next item returned is now 1, so `consume` will return `false`. | |
271 | /// assert_eq!(iter.next_if(|&x| x == 0), None); | |
272 | /// // `next_if` saves the value of the next item if it was not equal to `expected`. | |
273 | /// assert_eq!(iter.next(), Some(1)); | |
274 | /// ``` | |
275 | /// | |
276 | /// Consume any number less than 10. | |
277 | /// ``` | |
fc512014 XL |
278 | /// let mut iter = (1..20).peekable(); |
279 | /// // Consume all numbers less than 10 | |
280 | /// while iter.next_if(|&x| x < 10).is_some() {} | |
281 | /// // The next value returned will be 10 | |
282 | /// assert_eq!(iter.next(), Some(10)); | |
283 | /// ``` | |
5869c6ff | 284 | #[stable(feature = "peekable_next_if", since = "1.51.0")] |
fc512014 XL |
285 | pub fn next_if(&mut self, func: impl FnOnce(&I::Item) -> bool) -> Option<I::Item> { |
286 | match self.next() { | |
287 | Some(matched) if func(&matched) => Some(matched), | |
288 | other => { | |
289 | // Since we called `self.next()`, we consumed `self.peeked`. | |
290 | assert!(self.peeked.is_none()); | |
291 | self.peeked = Some(other); | |
292 | None | |
293 | } | |
294 | } | |
295 | } | |
296 | ||
297 | /// Consume and return the next item if it is equal to `expected`. | |
298 | /// | |
299 | /// # Example | |
300 | /// Consume a number if it's equal to 0. | |
301 | /// ``` | |
fc512014 XL |
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 | /// ``` | |
5869c6ff | 310 | #[stable(feature = "peekable_next_if", since = "1.51.0")] |
fc512014 XL |
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 | } |