itertools/
peeking_take_while.rs

1use std::iter::Peekable;
2use crate::PutBack;
3#[cfg(feature = "use_alloc")]
4use crate::PutBackN;
5
6/// An iterator that allows peeking at an element before deciding to accept it.
7///
8/// See [`.peeking_take_while()`](crate::Itertools::peeking_take_while)
9/// for more information.
10///
11/// This is implemented by peeking adaptors like peekable and put back,
12/// but also by a few iterators that can be peeked natively, like the slice’s
13/// by reference iterator (`std::slice::Iter`).
14pub trait PeekingNext : Iterator {
15    /// Pass a reference to the next iterator element to the closure `accept`;
16    /// if `accept` returns true, return it as the next element,
17    /// else None.
18    fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
19        where Self: Sized,
20              F: FnOnce(&Self::Item) -> bool;
21}
22
23impl<'a, I> PeekingNext for &'a mut I
24    where I: PeekingNext,
25{
26    fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
27        where F: FnOnce(&Self::Item) -> bool
28    {
29        (*self).peeking_next(accept)
30    }
31}
32
33impl<I> PeekingNext for Peekable<I>
34    where I: Iterator,
35{
36    fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
37        where F: FnOnce(&Self::Item) -> bool
38    {
39        if let Some(r) = self.peek() {
40            if !accept(r) {
41                return None;
42            }
43        }
44        self.next()
45    }
46}
47
48impl<I> PeekingNext for PutBack<I>
49    where I: Iterator,
50{
51    fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
52        where F: FnOnce(&Self::Item) -> bool
53    {
54        if let Some(r) = self.next() {
55            if !accept(&r) {
56                self.put_back(r);
57                return None;
58            }
59            Some(r)
60        } else {
61            None
62        }
63    }
64}
65
66#[cfg(feature = "use_alloc")]
67impl<I> PeekingNext for PutBackN<I>
68    where I: Iterator,
69{
70    fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
71        where F: FnOnce(&Self::Item) -> bool
72    {
73        if let Some(r) = self.next() {
74            if !accept(&r) {
75                self.put_back(r);
76                return None;
77            }
78            Some(r)
79        } else {
80            None
81        }
82    }
83}
84
85/// An iterator adaptor that takes items while a closure returns `true`.
86///
87/// See [`.peeking_take_while()`](crate::Itertools::peeking_take_while)
88/// for more information.
89#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
90pub struct PeekingTakeWhile<'a, I: 'a, F>
91    where I: Iterator,
92{
93    iter: &'a mut I,
94    f: F,
95}
96
97impl<'a, I: 'a, F> std::fmt::Debug for PeekingTakeWhile<'a, I, F>
98where
99    I: Iterator + std::fmt::Debug,
100{
101    debug_fmt_fields!(PeekingTakeWhile, iter);
102}
103
104/// Create a `PeekingTakeWhile`
105pub fn peeking_take_while<I, F>(iter: &mut I, f: F) -> PeekingTakeWhile<I, F>
106    where I: Iterator,
107{
108    PeekingTakeWhile {
109        iter,
110        f,
111    }
112}
113
114impl<'a, I, F> Iterator for PeekingTakeWhile<'a, I, F>
115    where I: PeekingNext,
116          F: FnMut(&I::Item) -> bool,
117
118{
119    type Item = I::Item;
120    fn next(&mut self) -> Option<Self::Item> {
121        self.iter.peeking_next(&mut self.f)
122    }
123
124    fn size_hint(&self) -> (usize, Option<usize>) {
125        (0, self.iter.size_hint().1)
126    }
127}
128
129impl<'a, I, F> PeekingNext for PeekingTakeWhile<'a, I, F>
130    where I: PeekingNext,
131          F: FnMut(&I::Item) -> bool,
132{
133    fn peeking_next<G>(&mut self, g: G) -> Option<Self::Item>
134        where G: FnOnce(&Self::Item) -> bool,
135    {
136        let f = &mut self.f;
137        self.iter.peeking_next(|r| f(r) && g(r))
138    }
139}
140
141// Some iterators are so lightweight we can simply clone them to save their
142// state and use that for peeking.
143macro_rules! peeking_next_by_clone {
144    ([$($typarm:tt)*] $type_:ty) => {
145        impl<$($typarm)*> PeekingNext for $type_ {
146            fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
147                where F: FnOnce(&Self::Item) -> bool
148            {
149                let saved_state = self.clone();
150                if let Some(r) = self.next() {
151                    if !accept(&r) {
152                        *self = saved_state;
153                    } else {
154                        return Some(r)
155                    }
156                }
157                None
158            }
159        }
160    }
161}
162
163peeking_next_by_clone! { ['a, T] ::std::slice::Iter<'a, T> }
164peeking_next_by_clone! { ['a] ::std::str::Chars<'a> }
165peeking_next_by_clone! { ['a] ::std::str::CharIndices<'a> }
166peeking_next_by_clone! { ['a] ::std::str::Bytes<'a> }
167peeking_next_by_clone! { ['a, T] ::std::option::Iter<'a, T> }
168peeking_next_by_clone! { ['a, T] ::std::result::Iter<'a, T> }
169peeking_next_by_clone! { [T] ::std::iter::Empty<T> }
170#[cfg(feature = "use_alloc")]
171peeking_next_by_clone! { ['a, T] alloc::collections::linked_list::Iter<'a, T> }
172#[cfg(feature = "use_alloc")]
173peeking_next_by_clone! { ['a, T] alloc::collections::vec_deque::Iter<'a, T> }
174
175// cloning a Rev has no extra overhead; peekable and put backs are never DEI.
176peeking_next_by_clone! { [I: Clone + PeekingNext + DoubleEndedIterator]
177                         ::std::iter::Rev<I> }