Flatten a 2d array into an iterator. An exercise for Iterator/IntoIterator trait bounds

flatten.rs

pub trait IteratorExt: Iterator
where
    Self: Sized,
    Self::Item: IntoIterator,
{
    fn flatten2(self) -> Flatten<Self> {
        return flatten(self);
    }
}

impl<T> IteratorExt for T
where
    T: Iterator,
    T::Item: IntoIterator,
{
}

pub fn flatten<O>(iter: O) -> Flatten<O::IntoIter>
where
    O: IntoIterator,
    O::Item: IntoIterator,
{
    return Flatten::new(iter.into_iter());
}

pub struct Flatten<O>
where
    O: Iterator,
    O::Item: IntoIterator,
{
    outer: O,
    front: Option<<O::Item as IntoIterator>::IntoIter>,
    back: Option<<O::Item as IntoIterator>::IntoIter>,
}

impl<O> Flatten<O>
where
    O: Iterator,
    O::Item: IntoIterator,
{
    fn new(iter: O) -> Self {
        return Self {
            outer: iter,
            front: None,
            back: None,
        };
    }
}

impl<O> Iterator for Flatten<O>
where
    O: Iterator,
    O::Item: IntoIterator,
{
    type Item = <O::Item as IntoIterator>::Item;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(ref mut iter_front) = self.front {
                if let Some(item) = iter_front.next() {
                    return Some(item);
                }
                self.front = None;
                continue;
            }

            if let Some(next_front) = self.outer.next() {
                self.front = Some(next_front.into_iter());
                continue;
            }
            return self.back.as_mut()?.next();
        }
    }
}

impl<O> DoubleEndedIterator for Flatten<O>
where
    O: DoubleEndedIterator,
    O::Item: IntoIterator,
    <O::Item as IntoIterator>::IntoIter: DoubleEndedIterator,
{
    fn next_back(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(ref mut iter_back) = self.back {
                if let Some(item) = iter_back.next_back() {
                    return Some(item);
                }
                self.back = None;
                continue;
            }

            if let Some(next_back) = self.outer.next_back() {
                self.back = Some(next_back.into_iter());
                continue;
            }
            return self.front.as_mut()?.next_back();
        }
    }
}

fn main() {
    println!("Hello, world!");
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn empty() {
        let mut it = flatten(std::iter::empty::<Vec<()>>());
        assert_eq!(it.next(), None);
    }
    #[test]
    fn empty_ext() {
        let it = std::iter::empty::<Vec<()>>();
        let mut it = it.flatten2();
        assert_eq!(it.next(), None);
    }

    #[test]
    fn two_wide() {
        let mut it = flatten(vec![vec![1, 2], vec![3, 4]]);
        assert_eq!(it.next(), Some(1));
        assert_eq!(it.next(), Some(2));
        assert_eq!(it.next(), Some(3));
        assert_eq!(it.next(), Some(4));
        assert_eq!(it.next(), None);
    }

    #[test]
    fn infinite_wide() {
        let mut it = flatten((0..).map(|i| 0..i));
        assert_eq!(it.next(), Some(0));
        assert_eq!(it.next(), Some(0));
        assert_eq!(it.next(), Some(1));
        assert_eq!(it.next(), Some(0));
        assert_eq!(it.next(), Some(1));
        assert_eq!(it.next(), Some(2));
    }

    #[test]
    fn two_wide_back() {
        let mut it = flatten(vec![vec![1, 2], vec![3, 4]]);
        assert_eq!(it.next(), Some(1));
        assert_eq!(it.next_back(), Some(4));
        assert_eq!(it.next(), Some(2));
        assert_eq!(it.next(), Some(3));
        assert_eq!(it.next(), None);
    }
}