take_until/
take_until.rs

1use core::fmt;
2use core::iter::FusedIterator;
3
4/// TakeUntilExt is an extension trait for iterators.
5/// It adds the [`Self::take_until`] method.
6pub trait TakeUntilExt<P>
7where
8    Self: Sized,
9{
10    /// Accepts elements until a predicate is true. Using this iterator is
11    /// equivalent to an **inclusive** [`Iterator::take_while`] with a negated
12    /// condition.
13    ///
14    /// # Example
15    ///
16    /// ## Parsing the next base 128 varint from a byte slice.
17    ///
18    /// ```rust
19    /// use take_until::TakeUntilExt;
20    ///
21    /// let varint = &[0b1010_1100u8, 0b0000_0010, 0b1000_0001];
22    /// let int: u32 = varint
23    ///     .iter()
24    ///     .take_until(|b| (**b & 0b1000_0000) == 0)
25    ///     .enumerate()
26    ///     .fold(0, |acc, (i, b)| {
27    ///         acc | ((*b & 0b0111_1111) as u32) << (i * 7)
28    ///      });
29    /// assert_eq!(300, int);
30    /// ```
31    ///
32    /// ## Take Until vs Take While (from Standard Library)
33    /// ```rust
34    /// use take_until::TakeUntilExt;
35    ///
36    /// let items = [1, 2, 3, 4, -5, -6, -7, -8];
37    /// let filtered_take_while = items
38    ///     .into_iter()
39    ///     .take_while(|x| *x > 0)
40    ///     .collect::<Vec<i32>>();
41    /// let filtered_take_until = items
42    ///     .into_iter()
43    ///     .take_until(|x| *x <= 0)
44    ///     .collect::<Vec<i32>>();
45    /// assert_eq!([1, 2, 3, 4], filtered_take_while.as_slice());
46    /// assert_eq!([1, 2, 3, 4, -5], filtered_take_until.as_slice());
47    /// ```
48    fn take_until(self, predicate: P) -> TakeUntil<Self, P>;
49}
50
51impl<I, P> TakeUntilExt<P> for I
52where
53    I: Sized + Iterator,
54    P: FnMut(&I::Item) -> bool,
55{
56    fn take_until(self, predicate: P) -> TakeUntil<Self, P> {
57        TakeUntil {
58            iter: self,
59            flag: false,
60            predicate,
61        }
62    }
63}
64/// TakeUntil is similar to the TakeWhile iterator,
65/// but takes items until the predicate is true,
66/// including the item that made the predicate true.
67pub struct TakeUntil<I, P> {
68    iter: I,
69    flag: bool,
70    predicate: P,
71}
72
73impl<I: fmt::Debug, P> fmt::Debug for TakeUntil<I, P> {
74    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
75        f.debug_struct("TakeUntil")
76            .field("iter", &self.iter)
77            .field("flag", &self.flag)
78            .finish()
79    }
80}
81
82impl<I, P> Iterator for TakeUntil<I, P>
83where
84    I: Iterator,
85    P: FnMut(&I::Item) -> bool,
86{
87    type Item = I::Item;
88
89    fn next(&mut self) -> Option<I::Item> {
90        if self.flag {
91            None
92        } else {
93            self.iter.next().map(|x| {
94                if (self.predicate)(&x) {
95                    self.flag = true;
96                }
97                x
98            })
99        }
100    }
101
102    fn size_hint(&self) -> (usize, Option<usize>) {
103        if self.flag {
104            (0, Some(0))
105        } else {
106            let (_, upper) = self.iter.size_hint();
107            (0, upper) // can't know a lower bound, due to the predicate
108        }
109    }
110}
111
112impl<I, P> FusedIterator for TakeUntil<I, P>
113where
114    I: FusedIterator,
115    P: FnMut(&I::Item) -> bool,
116{
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122    #[test]
123    fn test_size_hint_zero() {
124        let v: Vec<u8> = vec![0, 1, 2];
125        let mut iter = v.iter().take_until(|_| true);
126        assert_eq!((0, Some(3)), iter.size_hint());
127        iter.next();
128        assert_eq!((0, Some(0)), iter.size_hint());
129    }
130}