itertools/
peeking_take_while.rs

1use crate::PutBack;
2#[cfg(feature = "use_alloc")]
3use crate::PutBackN;
4use crate::RepeatN;
5use std::iter::Peekable;
6
7/// An iterator that allows peeking at an element before deciding to accept it.
8///
9/// See [`.peeking_take_while()`](crate::Itertools::peeking_take_while)
10/// for more information.
11///
12/// This is implemented by peeking adaptors like peekable and put back,
13/// but also by a few iterators that can be peeked natively, like the slice’s
14/// by reference iterator ([`std::slice::Iter`]).
15pub trait PeekingNext: Iterator {
16    /// Pass a reference to the next iterator element to the closure `accept`;
17    /// if `accept` returns `true`, return it as the next element,
18    /// else `None`.
19    fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
20    where
21        Self: Sized,
22        F: FnOnce(&Self::Item) -> bool;
23}
24
25impl<I> PeekingNext for &mut I
26where
27    I: PeekingNext,
28{
29    fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
30    where
31        F: FnOnce(&Self::Item) -> bool,
32    {
33        (*self).peeking_next(accept)
34    }
35}
36
37impl<I> PeekingNext for Peekable<I>
38where
39    I: Iterator,
40{
41    fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
42    where
43        F: FnOnce(&Self::Item) -> bool,
44    {
45        if let Some(r) = self.peek() {
46            if !accept(r) {
47                return None;
48            }
49        }
50        self.next()
51    }
52}
53
54impl<I> PeekingNext for PutBack<I>
55where
56    I: Iterator,
57{
58    fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
59    where
60        F: FnOnce(&Self::Item) -> bool,
61    {
62        if let Some(r) = self.next() {
63            if !accept(&r) {
64                self.put_back(r);
65                return None;
66            }
67            Some(r)
68        } else {
69            None
70        }
71    }
72}
73
74#[cfg(feature = "use_alloc")]
75impl<I> PeekingNext for PutBackN<I>
76where
77    I: Iterator,
78{
79    fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
80    where
81        F: FnOnce(&Self::Item) -> bool,
82    {
83        if let Some(r) = self.next() {
84            if !accept(&r) {
85                self.put_back(r);
86                return None;
87            }
88            Some(r)
89        } else {
90            None
91        }
92    }
93}
94
95#[cfg(feature = "use_alloc")]
96impl<T> PeekingNext for ::alloc::vec::IntoIter<T> {
97    fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
98    where
99        F: FnOnce(&Self::Item) -> bool,
100    {
101        match accept(self.as_slice().first()?) {
102            true => self.next(),
103            false => None,
104        }
105    }
106}
107
108#[cfg(feature = "use_alloc")]
109impl<'a, T> PeekingNext for ::alloc::vec::Drain<'a, T> {
110    fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
111    where
112        F: FnOnce(&Self::Item) -> bool,
113    {
114        match accept(self.as_slice().first()?) {
115            true => self.next(),
116            false => None,
117        }
118    }
119}
120
121#[cfg(feature = "use_alloc")]
122impl<'a> PeekingNext for ::alloc::string::Drain<'a> {
123    fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
124    where
125        F: FnOnce(&Self::Item) -> bool,
126    {
127        match accept(&self.as_str().chars().next()?) {
128            true => self.next(),
129            false => None,
130        }
131    }
132}
133
134impl<T, const N: usize> PeekingNext for ::core::array::IntoIter<T, N> {
135    fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
136    where
137        F: FnOnce(&Self::Item) -> bool,
138    {
139        match accept(self.as_slice().first()?) {
140            true => self.next(),
141            false => None,
142        }
143    }
144}
145
146impl<T: Clone> PeekingNext for RepeatN<T> {
147    fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
148    where
149        F: FnOnce(&Self::Item) -> bool,
150    {
151        let r = self.elt.as_ref()?;
152        if !accept(r) {
153            return None;
154        }
155        self.next()
156    }
157}
158
159/// An iterator adaptor that takes items while a closure returns `true`.
160///
161/// See [`.peeking_take_while()`](crate::Itertools::peeking_take_while)
162/// for more information.
163#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
164pub struct PeekingTakeWhile<'a, I, F>
165where
166    I: Iterator + 'a,
167{
168    iter: &'a mut I,
169    f: F,
170}
171
172impl<'a, I, F> std::fmt::Debug for PeekingTakeWhile<'a, I, F>
173where
174    I: Iterator + std::fmt::Debug + 'a,
175{
176    debug_fmt_fields!(PeekingTakeWhile, iter);
177}
178
179/// Create a `PeekingTakeWhile`
180pub fn peeking_take_while<I, F>(iter: &mut I, f: F) -> PeekingTakeWhile<'_, I, F>
181where
182    I: Iterator,
183{
184    PeekingTakeWhile { iter, f }
185}
186
187impl<I, F> Iterator for PeekingTakeWhile<'_, I, F>
188where
189    I: PeekingNext,
190    F: FnMut(&I::Item) -> bool,
191{
192    type Item = I::Item;
193    fn next(&mut self) -> Option<Self::Item> {
194        self.iter.peeking_next(&mut self.f)
195    }
196
197    fn size_hint(&self) -> (usize, Option<usize>) {
198        (0, self.iter.size_hint().1)
199    }
200}
201
202impl<I, F> PeekingNext for PeekingTakeWhile<'_, I, F>
203where
204    I: PeekingNext,
205    F: FnMut(&I::Item) -> bool,
206{
207    fn peeking_next<G>(&mut self, g: G) -> Option<Self::Item>
208    where
209        G: FnOnce(&Self::Item) -> bool,
210    {
211        let f = &mut self.f;
212        self.iter.peeking_next(|r| f(r) && g(r))
213    }
214}
215
216// Some iterators are so lightweight we can simply clone them to save their
217// state and use that for peeking.
218macro_rules! peeking_next_by_clone {
219    ([$($typarm:tt)*] $type_:ty) => {
220        impl<$($typarm)*> PeekingNext for $type_ {
221            fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
222                where F: FnOnce(&Self::Item) -> bool
223            {
224                let saved_state = self.clone();
225                if let Some(r) = self.next() {
226                    if !accept(&r) {
227                        *self = saved_state;
228                    } else {
229                        return Some(r)
230                    }
231                }
232                None
233            }
234        }
235    }
236}
237
238peeking_next_by_clone! { ['a, T] ::std::slice::Iter<'a, T> }
239peeking_next_by_clone! { ['a] ::std::str::Chars<'a> }
240peeking_next_by_clone! { ['a] ::std::str::CharIndices<'a> }
241peeking_next_by_clone! { ['a] ::std::str::Bytes<'a> }
242peeking_next_by_clone! { ['a, T] ::std::option::Iter<'a, T> }
243peeking_next_by_clone! { ['a, T] ::std::result::Iter<'a, T> }
244peeking_next_by_clone! { [T] ::std::iter::Empty<T> }
245
246#[cfg(feature = "use_alloc")]
247peeking_next_by_clone! { ['a, T] alloc::collections::linked_list::Iter<'a, T> }
248#[cfg(feature = "use_alloc")]
249peeking_next_by_clone! { ['a, T] alloc::collections::vec_deque::Iter<'a, T> }
250
251#[cfg(feature = "use_alloc")]
252peeking_next_by_clone! { ['a, K, V] alloc::collections::btree_map::Iter<'a, K, V> }
253#[cfg(feature = "use_alloc")]
254peeking_next_by_clone! { ['a, K, V] alloc::collections::btree_map::Keys<'a, K, V> }
255#[cfg(feature = "use_alloc")]
256peeking_next_by_clone! { ['a, K, V] alloc::collections::btree_map::Values<'a, K, V> }
257
258#[cfg(feature = "use_alloc")]
259peeking_next_by_clone! { ['a, T] alloc::collections::btree_set::Iter<'a, T> }
260#[cfg(feature = "use_alloc")]
261peeking_next_by_clone! { ['a, T] alloc::collections::binary_heap::Iter<'a, T> }
262
263#[cfg(feature = "use_std")]
264peeking_next_by_clone! { ['a, K, V] std::collections::hash_map::Iter<'a, K, V> }
265#[cfg(feature = "use_std")]
266peeking_next_by_clone! { ['a, K, V] std::collections::hash_map::Keys<'a, K, V> }
267#[cfg(feature = "use_std")]
268peeking_next_by_clone! { ['a, K, V] std::collections::hash_map::Values<'a, K, V> }
269#[cfg(feature = "use_std")]
270peeking_next_by_clone! { ['a, T] std::collections::hash_set::Iter<'a, T> }
271
272// cloning a Rev has no extra overhead; peekable and put backs are never DEI.
273peeking_next_by_clone! { [I: Clone + PeekingNext + DoubleEndedIterator]
274::std::iter::Rev<I> }
itertools/peeking_take_while.rs

itertools/
peeking_take_while.rs
