`chain_with` extension method to `Iterator` that evaluate chained iterator lazily

chain_with.rs

pub trait IterExt {
    type Item;
    fn chain_with<F, NI>(self, next_iter_fn: F) -> ChainedWithIter<Self, F, NI>
    where
        Self: Sized,
        NI: Iterator<Item = Self::Item> + Sized,
        F: FnMut() -> NI + Sized,
    {
        ChainedWithIter::new(self, next_iter_fn)
    }
}

#[cfg(test)]
mod tests {
    use crate::util::iter::IterExt;
    use std::cell::RefCell;
    use std::rc::Rc;

    #[test]
    fn chained_side_effects_should_happen_after_first_iterator_done() {
        let mut side_effects: Rc<RefCell<Vec<_>>> = Default::default();
        let mut num: Rc<RefCell<_>> = Default::default();

        let iter_first = std::iter::repeat_with(|| {
            *num.borrow_mut() += 1;
            let new_num = *num.borrow();
            side_effects.borrow_mut().push(new_num);
            new_num
        })
        .take(3);

        let iter = iter_first.chain_with(|| {
            *num.borrow_mut() = 0;
            side_effects.borrow_mut().push(*num.borrow());
            std::iter::once(0).chain(
                std::iter::repeat_with(|| {
                    *num.borrow_mut() -= 1;
                    let new_num = *num.borrow();
                    side_effects.borrow_mut().push(new_num);
                    new_num
                })
                .take(3),
            )
        });
        let collected: Vec<_> = iter.collect();

        assert_eq!(*side_effects.borrow(), vec![1, 2, 3, 0, -1, -2, -3]);
        assert_eq!(*side_effects.borrow(), collected);
    }
}

impl<T> IterExt for T
where
    T: Iterator,
{
    type Item = <T as Iterator>::Item;
}

pub struct ChainedWithIter<I, F, NI> {
    current_iter: I,
    current_finished: bool,
    next_iter_fn: F,
    next_iter: Option<NI>,
}
impl<I, F, NI> ChainedWithIter<I, F, NI> {
    pub fn new(current_iter: I, next_iter_fn: F) -> Self {
        Self {
            current_iter,
            current_finished: false,
            next_iter_fn,
            next_iter: None,
        }
    }
}
impl<I, F, NI, E> Iterator for ChainedWithIter<I, F, NI>
where
    I: Iterator<Item = E>,
    NI: Iterator<Item = E>,
    F: FnMut() -> NI,
{
    type Item = E;
    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if self.current_finished {
                return self.next_iter.as_mut().map(|i| i.next()).flatten();
            }
            match self.current_iter.next() {
                r @ Some(_) => return r,
                None => {
                    self.current_finished = true;
                    self.next_iter = Some((self.next_iter_fn)());
                    continue;
                }
            }
        }
    }
}