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
95impl<T: Clone> PeekingNext for RepeatN<T> {
96    fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
97    where
98        F: FnOnce(&Self::Item) -> bool,
99    {
100        let r = self.elt.as_ref()?;
101        if !accept(r) {
102            return None;
103        }
104        self.next()
105    }
106}
107
108/// An iterator adaptor that takes items while a closure returns `true`.
109///
110/// See [`.peeking_take_while()`](crate::Itertools::peeking_take_while)
111/// for more information.
112#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
113pub struct PeekingTakeWhile<'a, I, F>
114where
115    I: Iterator + 'a,
116{
117    iter: &'a mut I,
118    f: F,
119}
120
121impl<'a, I, F> std::fmt::Debug for PeekingTakeWhile<'a, I, F>
122where
123    I: Iterator + std::fmt::Debug + 'a,
124{
125    debug_fmt_fields!(PeekingTakeWhile, iter);
126}
127
128/// Create a `PeekingTakeWhile`
129pub fn peeking_take_while<I, F>(iter: &mut I, f: F) -> PeekingTakeWhile<I, F>
130where
131    I: Iterator,
132{
133    PeekingTakeWhile { iter, f }
134}
135
136impl<I, F> Iterator for PeekingTakeWhile<'_, I, F>
137where
138    I: PeekingNext,
139    F: FnMut(&I::Item) -> bool,
140{
141    type Item = I::Item;
142    fn next(&mut self) -> Option<Self::Item> {
143        self.iter.peeking_next(&mut self.f)
144    }
145
146    fn size_hint(&self) -> (usize, Option<usize>) {
147        (0, self.iter.size_hint().1)
148    }
149}
150
151impl<I, F> PeekingNext for PeekingTakeWhile<'_, I, F>
152where
153    I: PeekingNext,
154    F: FnMut(&I::Item) -> bool,
155{
156    fn peeking_next<G>(&mut self, g: G) -> Option<Self::Item>
157    where
158        G: FnOnce(&Self::Item) -> bool,
159    {
160        let f = &mut self.f;
161        self.iter.peeking_next(|r| f(r) && g(r))
162    }
163}
164
165// Some iterators are so lightweight we can simply clone them to save their
166// state and use that for peeking.
167macro_rules! peeking_next_by_clone {
168    ([$($typarm:tt)*] $type_:ty) => {
169        impl<$($typarm)*> PeekingNext for $type_ {
170            fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
171                where F: FnOnce(&Self::Item) -> bool
172            {
173                let saved_state = self.clone();
174                if let Some(r) = self.next() {
175                    if !accept(&r) {
176                        *self = saved_state;
177                    } else {
178                        return Some(r)
179                    }
180                }
181                None
182            }
183        }
184    }
185}
186
187peeking_next_by_clone! { ['a, T] ::std::slice::Iter<'a, T> }
188peeking_next_by_clone! { ['a] ::std::str::Chars<'a> }
189peeking_next_by_clone! { ['a] ::std::str::CharIndices<'a> }
190peeking_next_by_clone! { ['a] ::std::str::Bytes<'a> }
191peeking_next_by_clone! { ['a, T] ::std::option::Iter<'a, T> }
192peeking_next_by_clone! { ['a, T] ::std::result::Iter<'a, T> }
193peeking_next_by_clone! { [T] ::std::iter::Empty<T> }
194#[cfg(feature = "use_alloc")]
195peeking_next_by_clone! { ['a, T] alloc::collections::linked_list::Iter<'a, T> }
196#[cfg(feature = "use_alloc")]
197peeking_next_by_clone! { ['a, T] alloc::collections::vec_deque::Iter<'a, T> }
198
199// cloning a Rev has no extra overhead; peekable and put backs are never DEI.
200peeking_next_by_clone! { [I: Clone + PeekingNext + DoubleEndedIterator]
201::std::iter::Rev<I> }
itertools/peeking_take_while.rs

itertools/
peeking_take_while.rs
