1//! Provides the `peeking_take_while` iterator adaptor method.
2//!
3//! The `peeking_take_while` method is very similar to `take_while`, but behaves
4//! differently when used with a borrowed iterator (perhaps returned by
5//! `Iterator::by_ref`).
6//!
7//! `peeking_take_while` peeks at the next item in the iterator and runs the
8//! predicate on that peeked item. This avoids consuming the first item yielded
9//! by the underlying iterator for which the predicate returns `false`. On the
10//! other hand, `take_while` will consume that first item for which the
11//! predicate returns `false`, and it will be lost.
12//!
13//! In case the closure may have side effects, it could be necessary to apply
14//! [`fuse`](Iterator::fuse) on the returned iterator, to prevent the predicate
15//! from being called after it first returned `false`.
16//!
17//! ```
18//! // Bring the `peeking_take_while` method for peekable iterators into
19//! // scope.
20//! use peeking_take_while::PeekableExt;
21//!
22//! # fn main() {
23//! // Let's say we have two collections we want to iterate through: `xs` and
24//! // `ys`. We want to perform one operation on all the leading contiguous
25//! // elements that match some predicate, and a different thing with the rest of
26//! // the elements. With the `xs`, we will use the normal `take_while`. With the
27//! // `ys`, we will use `peeking_take_while`.
28//!
29//! let xs: Vec<u8> = (0..100).collect();
30//! let ys = xs.clone();
31//!
32//! let mut iter_xs = xs.into_iter();
33//! let mut iter_ys = ys.into_iter().peekable();
34//!
35//! {
36//!     // Let's do one thing with all the items that are less than 10.
37//! #   fn do_things_with<T>(_: T) {}
38//!
39//!     let xs_less_than_ten = iter_xs.by_ref().take_while(|x| *x < 10);
40//!     for x in xs_less_than_ten {
41//!         do_things_with(x);
42//!     }
43//!
44//!     let ys_less_than_ten = iter_ys.by_ref().peeking_take_while(|y| *y < 10);
45//!     for y in ys_less_than_ten {
46//!         do_things_with(y);
47//!     }
48//! }
49//!
50//! // And now we will do some other thing with the items that are greater than
51//! // or equal to 10.
52//!
53//! // ...except, when using plain old `take_while` we lost 10!
54//! assert_eq!(iter_xs.next(), Some(11));
55//!
56//! // However, when using `peeking_take_while` we did not! Great!
57//! assert_eq!(iter_ys.next(), Some(10));
58//! # }
59//! ```
60
61#![no_std]
62#![forbid(
63    clippy::as_conversions,
64    clippy::cast_ptr_alignment,
65    missing_docs,
66    trivial_casts,
67    unsafe_code
68)]
69
70use core::fmt;
71
72/// The `Iterator` extension trait that provides the `peeking_take_while`
73/// method.
74///
75/// See the [module documentation](./index.html) for details.
76pub trait PeekableExt<I>: Iterator
77where
78    I: Iterator,
79{
80    /// The `peeking_take_while` method is very similar to `take_while`, but behaves
81    /// differently when used with a borrowed iterator (perhaps returned by
82    /// `Iterator::by_ref`).
83    ///
84    /// `peeking_take_while` peeks at the next item in the iterator and runs the
85    /// predicate on that peeked item. This avoids consuming the first item yielded
86    /// by the underlying iterator for which the predicate returns `false`. On the
87    /// other hand, `take_while` will consume that first item for which the
88    /// predicate returns `false`, and it will be lost.
89    ///
90    /// In contrast to `take_while`, iterating the iterator might call the predicate again
91    /// after it first returned `false` (the returned iterator isn't fused).
92    /// If that is not intended, calling [`fuse`](Iterator::fuse) on the returned iterator
93    /// prevents that.
94    fn peeking_take_while<P>(&mut self, predicate: P) -> PeekingTakeWhile<'_, I, P>
95    where
96        P: FnMut(&Self::Item) -> bool;
97}
98
99impl<I: Iterator> PeekableExt<I> for core::iter::Peekable<I> {
100    #[inline]
101    fn peeking_take_while<P>(&mut self, predicate: P) -> PeekingTakeWhile<'_, I, P>
102    where
103        P: FnMut(&Self::Item) -> bool,
104    {
105        PeekingTakeWhile {
106            iter: self,
107            predicate,
108        }
109    }
110}
111
112/// The iterator returned by `peeking_take_while`.
113///
114/// See the [module documentation](./index.html) for details.
115pub struct PeekingTakeWhile<'a, I, P>
116where
117    I: Iterator,
118{
119    pub(crate) iter: &'a mut core::iter::Peekable<I>,
120    pub(crate) predicate: P,
121}
122
123impl<I, P> fmt::Debug for PeekingTakeWhile<'_, I, P>
124where
125    I: Iterator + fmt::Debug,
126    I::Item: fmt::Debug,
127{
128    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
129        f.debug_struct("PeekingTakeWhile")
130            .field("iter", &self.iter)
131            .finish()
132    }
133}
134
135impl<I, P> Iterator for PeekingTakeWhile<'_, I, P>
136where
137    I: Iterator,
138    P: FnMut(&I::Item) -> bool,
139{
140    type Item = I::Item;
141
142    #[inline]
143    fn next(&mut self) -> Option<Self::Item> {
144        self.iter.next_if(&mut self.predicate)
145    }
146
147    #[inline]
148    fn size_hint(&self) -> (usize, Option<usize>) {
149        // can't know a lower bound, due to the predicate
150        (0, self.iter.size_hint().1)
151    }
152
153    #[inline]
154    fn fold<B, F>(mut self, mut accum: B, mut f: F) -> B
155    where
156        F: FnMut(B, I::Item) -> B,
157    {
158        while let Some(x) = self.iter.next_if(&mut self.predicate) {
159            accum = f(accum, x);
160        }
161        accum
162    }
163}
164
165// interestingly, `PeekingTakeWhile` is not automatically fused,
166// even when the inner iterator is fused, see also: `tests::not_fused`.
167
168#[cfg(test)]
169mod tests {
170    use crate::PeekableExt;
171
172    #[test]
173    fn basic() {
174        let mut it0 = (1..11).peekable();
175        let a: u32 = it0.peeking_take_while(|&i| i < 5).sum();
176        let b: u32 = it0.sum();
177        assert_eq!(a, 10);
178        assert_eq!(b, 45);
179    }
180
181    #[test]
182    fn basic_fused() {
183        let mut it0 = (1..11).peekable();
184        let a: u32 = it0.peeking_take_while(|&i| i < 5).fuse().sum();
185        let b: u32 = it0.sum();
186        assert_eq!(a, 10);
187        assert_eq!(b, 45);
188    }
189
190    #[test]
191    fn not_fused() {
192        let mut it0 = (0..10).peekable();
193        let mut ax = true;
194        let mut it1 = it0.peeking_take_while(|_| {
195            ax = !ax;
196            ax
197        });
198        assert!(it1.next().is_none());
199        assert_eq!(it1.next(), Some(0));
200        assert!(it1.next().is_none());
201        assert_eq!(it1.next(), Some(1));
202        assert_eq!(ax, true);
203    }
204
205    #[test]
206    fn fused() {
207        let mut it0 = (0..10).peekable();
208        let mut ax = true;
209        let mut it1 = it0
210            .peeking_take_while(|_| {
211                ax = !ax;
212                ax
213            })
214            .fuse();
215        assert!(it1.next().is_none());
216        assert!(it1.next().is_none());
217        assert_eq!(ax, false);
218    }
219}
220