Rust Cartesian Product Iterator

main.rs

mod util;
use std::ops::Range;
use util::*;

fn main() {
    let v: Vec<Range<i32>> = [(1..4), (11..14), (111..114)]
        .into_iter()
        .map(|r| r.clone())
        .collect();
    for items in v.into_iter().flat_cprod() {
        println!("{:?}", items);
    }
}

output.txt

[1, 11, 111]
[1, 11, 112]
[1, 11, 113]
[1, 12, 111]
[1, 12, 112]
[1, 12, 113]
[1, 13, 111]
[1, 13, 112]
[1, 13, 113]
[2, 11, 111]
[2, 11, 112]
[2, 11, 113]
[2, 12, 111]
[2, 12, 112]
[2, 12, 113]
[2, 13, 111]
[2, 13, 112]
[2, 13, 113]
[3, 11, 111]
[3, 11, 112]
[3, 11, 113]
[3, 12, 111]
[3, 12, 112]
[3, 12, 113]
[3, 13, 111]
[3, 13, 112]
[3, 13, 113]

util.rs

use std::iter::Peekable;

// Followed cartesian_product from Itertools
pub trait IteratorUtil: Iterator {
    fn flat_cprod<J, I>(self) -> FlatCProd<J, I>
    where
        Self: Sized,
        Self::Item: Clone + IntoIterator<Item = I, IntoIter = J>,
        J: Clone + Iterator<Item = I>,
        I: Clone,
    {
        let initial_columns: Vec<J> = self.map(|i| i.into_iter()).collect();
        let columns = initial_columns
            .clone()
            .into_iter()
            .map(|i| i.peekable())
            .collect();
        FlatCProd {
            columns,
            initial_columns,
        }
    }
}

#[derive(Debug, Clone)]
pub struct FlatCProd<J, I>
where
    J: Clone + Iterator<Item = I>,
    I: Clone,
{
    initial_columns: Vec<J>,
    columns: Vec<Peekable<J>>,
}

impl<J: Clone + Iterator<Item = I>, I: Clone> FlatCProd<J, I> {
    fn peek_next(&mut self) -> Option<Vec<I>> {
        return self
            .columns
            .iter_mut()
            .map(|i| i.peek().map(|o| o.clone()))
            .collect();
    }
}

impl<J: Clone + Iterator<Item = I>, I: Clone> Iterator for FlatCProd<J, I> {
    type Item = Vec<I>;
    fn next(&mut self) -> Option<Self::Item> {
        if self.columns.len() == 0 {
            return None;
        }
        let next = self.peek_next();
        match next {
            None => return None,
            Some(curr) => {
                'l: for i in (0..self.columns.len()).rev() {
                    // This is always Some due to peek_next
                    self.columns[i].next();
                    match self.columns[i].peek() {
                        Some(_) => break 'l,
                        None => {
                            // Reset the column if it's not the last one
                            if i > 0 {
                                self.columns[i] = self.initial_columns[i].clone().peekable();
                            }
                        }
                    };
                }
                return Some(curr);
            }
        }
    }
}

impl<I> IteratorUtil for I where I: Iterator {}