Flatten a nested Iterator in Rust

flatten.rs

// Author: Victor I. Afolabi
//
// Exploring Iterator Traits in Rust to implement Flatten functionality of 2D-iterators.
//

pub trait IteratorExt: Iterator {
    fn flatten_ext(self) -> Flatten<Self>
    where
        Self: Sized + std::fmt::Debug,
        Self::Item: IntoIterator,
        <Self::Item as IntoIterator>::IntoIter: std::fmt::Debug;
}

impl<T> IteratorExt for T
where
    T: Iterator,
{
    fn flatten_ext(self) -> Flatten<Self>
    where
        Self: Sized + std::fmt::Debug,
        Self::Item: IntoIterator,
        <Self::Item as IntoIterator>::IntoIter: std::fmt::Debug,
    {
        flatten(self)
    }
}

/// Flatten a 2-D nested Iterator. O should be convertable `.into_iter()` not neccessarily
/// an `iter()` itself.
pub fn flatten<I>(iter: I) -> Flatten<I::IntoIter>
where
    I: IntoIterator, // Type should be convertable into an iterator, not neccessarily an iterator itself.
    <I as IntoIterator>::Item: IntoIterator,
    <I::Item as IntoIterator>::IntoIter: std::fmt::Debug,
{
    Flatten::new(iter.into_iter())
}

#[derive(Debug)]
pub struct Flatten<O>
where
    O: Iterator,
    <O as Iterator>::Item: IntoIterator,
    <O::Item as IntoIterator>::IntoIter: std::fmt::Debug,
{
    outer: O,
    // inner: Option<<O::Item as IntoIterator>::IntoIter>,
    front_iter: Option<<O::Item as IntoIterator>::IntoIter>,
    back_iter: Option<<O::Item as IntoIterator>::IntoIter>,
}

impl<O> Flatten<O>
where
    O: Iterator,
    <O as Iterator>::Item: IntoIterator,
    <O::Item as IntoIterator>::IntoIter: std::fmt::Debug,
{
    fn new(iter: O) -> Self {
        Flatten {
            outer: iter,
            // inner: None,
            front_iter: None,
            back_iter: None,
        }
    }
}

impl<O> Iterator for Flatten<O>
where
    O: Iterator,                         // Outer iteraotor implements Iterator.
    <O as Iterator>::Item: IntoIterator, // Outer iterator's item implements IntoIterator
    <O::Item as IntoIterator>::IntoIter: std::fmt::Debug,
{
    // Yeilds the inner type.
    type Item = <O::Item as IntoIterator>::Item;

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

            if let Some(next_inner) = self.outer.next() {
                self.front_iter = Some(next_inner.into_iter());
            } else {
                return self.back_iter.as_mut()?.next();
            }
        }
    }
}

impl<O> DoubleEndedIterator for Flatten<O>
where
    O: Iterator + DoubleEndedIterator,
    <O as Iterator>::Item: IntoIterator,
    <O::Item as IntoIterator>::IntoIter: DoubleEndedIterator,
    <O::Item as IntoIterator>::IntoIter: std::fmt::Debug,
{
    fn next_back(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(ref mut back_iter) = self.back_iter {
                if let Some(i) = back_iter.next_back() {
                    return Some(i);
                }
                self.back_iter = None;
            }

            if let Some(next_inner) = self.outer.next_back() {
                self.back_iter = Some(next_inner.into_iter());
            } else {
                return self.front_iter.as_mut()?.next_back();
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn empty() {
        // Size.
        assert_eq!(flatten(std::iter::empty::<Vec<()>>()).count(), 0);

        // Values.
        assert_eq!(
            flatten(std::iter::empty::<Vec<()>>()).collect::<Vec<()>>(),
            &[]
        );
    }

    #[test]
    fn empty_wide() {
        // Size.
        assert_eq!(flatten(vec![Vec::<()>::new(), vec![], vec![]]).count(), 0);

        // Values.
        assert_eq!(
            flatten(vec![Vec::<()>::new(), vec![], vec![]]).collect::<Vec<()>>(),
            &[]
        );
    }

    #[test]
    fn one() {
        // Size.
        assert_eq!(flatten(std::iter::once(vec!["a"])).count(), 1);

        // Values.
        assert_eq!(
            flatten(std::iter::once(vec!["a"])).collect::<Vec<&str>>(),
            &["a"]
        );
    }

    #[test]
    fn two() {
        // Size.
        assert_eq!(flatten(std::iter::once(vec!["a", "b"])).count(), 2);

        // Values.
        assert_eq!(
            flatten(std::iter::once(vec!["a", "b"])).collect::<Vec<&str>>(),
            &["a", "b"]
        );
    }

    #[test]
    fn two_wide() {
        // size.
        assert_eq!(flatten(vec![vec!["a"], vec!["b"]]).count(), 2);

        // Values.
        assert_eq!(
            flatten(vec![vec!["a"], vec!["b"]]).collect::<Vec<&str>>(),
            &["a", "b"]
        );
    }

    #[test]
    fn reverse_two() {
        // Size.
        assert_eq!(flatten(std::iter::once(vec!["a", "b"])).rev().count(), 2);

        // Values.
        assert_eq!(
            flatten(std::iter::once(vec!["a", "b"]))
                .rev()
                .collect::<Vec<&str>>(),
            &["b", "a"]
        );
    }

    #[test]
    fn reverse_two_wide() {
        // size.
        assert_eq!(flatten(vec![vec!["a"], vec!["b"]]).rev().count(), 2);

        // Values.
        assert_eq!(
            flatten(vec![vec!["a"], vec!["b"]])
                .rev()
                .collect::<Vec<&str>>(),
            &["b", "a"]
        );
    }

    #[test]
    fn both_ends_str() {
        let mut iter = flatten(vec![vec!["a", "b"], vec!["c", "d"]]);
        assert_eq!(iter.next(), Some("a"));
        assert_eq!(iter.next_back(), Some("d"));
        assert_eq!(iter.next(), Some("b"));
        assert_eq!(iter.next_back(), Some("c"));
        assert_eq!(iter.next(), None);
        assert_eq!(iter.next_back(), None);
    }

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

    #[test]
    fn inf() {
        let mut iter = flatten((0..).map(|i| 0..i));
        // 0 => 0..0 => empty
        // 1 => 0..1 => [0]
        // 2 => 0..2 => [0, 1]
        assert_eq!(iter.next(), Some(0));
        assert_eq!(iter.next(), Some(0));
        assert_eq!(iter.next(), Some(1));
    }

    #[test]
    fn deep() {
        assert_eq!(flatten(flatten(vec![vec![vec!["a", "b"]]])).count(), 2);

        // Values.
        assert_eq!(
            flatten(flatten(vec![vec![vec!["a", "b"]]])).collect::<Vec<&str>>(),
            &["a", "b"]
        );

        // Reverse.
        assert_eq!(
            flatten(flatten(vec![vec![vec!["a", "b"]]]))
                .rev()
                .collect::<Vec<&str>>(),
            &["b", "a"]
        );
    }

    #[test]
    fn ext() {
        assert_eq!(vec![vec![0, 1]].into_iter().flatten_ext().count(), 2);
    }
}