rust iterators: take_while_inclusive() - take_while including the first non-matching element

take_while_inclusive.rs

enum TakeWhileInclusive<I, P> {
    Taking(I, P),
    Done,
}

impl<I, P> Iterator for TakeWhileInclusive<I, P>
where
    P: FnMut(&I::Item) -> bool,
    I: Iterator,
{
    type Item = I::Item;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        match self {
            TakeWhileInclusive::Taking(iter, predicate) => {
                let next = iter.next();
                if !next.as_ref().map_or(false, predicate) {
                    *self = TakeWhileInclusive::Done;
                }
                next
            }
            TakeWhileInclusive::Done => None,
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        match self {
            TakeWhileInclusive::Taking(iter, _) => {
                (0, iter.size_hint().1)
            },
            TakeWhileInclusive::Done => (0, Some(0)),
        }
    }
}

trait IteratorExtenstion: Iterator + Sized {
    fn take_while_inclusive<P>(self, predicate: P) -> TakeWhileInclusive<Self, P>
    where
        P: FnMut(&Self::Item) -> bool,
    {
        TakeWhileInclusive::Taking(self, predicate)
    }
}

impl<T> IteratorExtenstion for T where T: Iterator {}

impl<I, P> std::iter::FusedIterator for TakeWhileInclusive<I, P>
where
    P: FnMut(&I::Item) -> bool,
    I: Iterator {

}

#[test]
fn test() {
    assert_eq!([1i32,2,3,4,5]
                .iter()
                .copied()
                .take_while_inclusive(|&x| x < 4)
                .collect::<Vec<_>>(), 
                vec![1i32, 2, 3, 4]);
                
    assert_eq!([1i32,2,3,4,5]
                .iter()
                .copied()
                .take_while_inclusive(|&x| x < 6)
                .collect::<Vec<_>>(), 
                vec![1i32, 2, 3, 4, 5]);
                
    assert_eq!([1i32,2,3,4,5]
                .iter()
                .copied()
                .take_while_inclusive(|_| false)
                .collect::<Vec<_>>(), 
                vec![1i32]);
                
    assert_eq!([0i32; 0]
                .iter()
                .copied()
                .take_while_inclusive(|&x| x < 4)
                .collect::<Vec<i32>>(), 
                vec![]);
}