rust-itertools / itertools

Extra iterator adaptors, iterator methods, free functions, and macros.
https://docs.rs/itertools/
Apache License 2.0
2.72k stars 309 forks source link

Iterator that omits after a certain count, providing a substitute #676

Open Enyium opened 1 year ago

Enyium commented 1 year ago

I wrote this code for my purposes. But you can feel free to edit it and incorparate it into your crate.

Example transformations

This is a test sentence.
 |
 v
This is a ...
Short.
 |
 v
Short.
here
you
see
many
lines
 |
 v
here
you
see
...
two
lines
 |
 v
two
lines

Code

mod omiter {
    //! Example usage:
    //!
    //! ```
    //! string.chars().omit_after(10, "...".chars());
    //! string.lines().omit_after(3, iter::once("..."));
    //! ```

    pub trait Omiter<I: IntoIterator> {
        fn omit_after<J>(self, n: usize, substitute: J) -> OmitAfter<I::IntoIter, J::IntoIter>
        where
            J: IntoIterator<Item = I::Item>;
    }

    impl<I: IntoIterator> Omiter<I> for I {
        fn omit_after<J>(self, n: usize, substitute: J) -> OmitAfter<I::IntoIter, J::IntoIter>
        where
            J: IntoIterator<Item = I::Item>,
        {
            OmitAfter {
                iter: Some(self.into_iter()),
                quota: n,
                substitute_iter: Some(substitute.into_iter()),
            }
        }
    }

    pub struct OmitAfter<I, J> {
        iter: Option<I>,
        quota: usize,
        substitute_iter: Option<J>,
    }

    impl<I, J> Iterator for OmitAfter<I, J>
    where
        I: Iterator,
        J: Iterator<Item = I::Item>,
    {
        type Item = <I as Iterator>::Item;

        fn next(&mut self) -> Option<<I as Iterator>::Item> {
            let item = if let Some(iter) = self.iter.as_mut() {
                if let Some(item) = iter.next() {
                    if self.quota > 0 {
                        self.quota -= 1;
                        Some(item)
                    } else {
                        self.iter = None;
                        None
                    }
                } else {
                    // Main iterator finished. Prevent feeding on substitute.
                    self.iter = None;
                    self.substitute_iter = None;
                    return None;
                }
            } else {
                None
            };

            if let (None, Some(substitute_iter)) = (&item, self.substitute_iter.as_mut()) {
                substitute_iter.next()
            } else {
                item
            }
        }
    }
}
phimuemue commented 1 year ago

Hi there, just to clarify: iter.omit_after(n, substitute_iter) is a shorthand for iter.take(n).chain(substitute_iter), right?

Enyium commented 1 year ago

No, the substitute iterator isn't fed off of when there's nothing to be omitted after n items.