Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

How to iterate over the product of several ranges or iterators?

Is there a natural way in Rust to iterate over the "product" of several ranges or iterators?

This comes up when you're iterating over a multidimensional array, or perhaps some state space. For instance, I want to consider all possible values of a boolean tuple with 5 elements. Nesting 5 for loops is a bit unwieldy.

like image 901
starwed Avatar asked Jul 17 '14 21:07

starwed


1 Answers

Here is a macro that does the job:

macro_rules! product {
    ($first:ident, $($next:ident),*) => (
        $first.iter() $(
            .flat_map(|e| std::iter::repeat(e)
                .zip($next.iter()))
        )*
    );
}

fn main() {
    let a = ['A', 'B', 'C'];
    let b = [1, 4];
    let c = [true, false];
    let d = ['x', 'y'];
    
    for (((a, b), c), d) in product![a, b, c, d] {
        println!("{} {} {} {}", a, b, c, d);
    }
}

Output:

A 1 true x
A 1 true y
A 1 false x
A 1 false y
A 4 true x
A 4 true y
etc...

Playpen example

The macro expands to the following

a.iter()
    .flat_map(|e| std::iter::repeat(e).zip(b.iter()))
    .flat_map(|e| std::iter::repeat(e).zip(c.iter()))
    .flat_map(|e| std::iter::repeat(e).zip(d.iter()))

flat_map(|e| ... ) combines a sequence of iterators into an iterator. The e is an element yielded by an iterator.

std::iter::repeat(e) creates an iterator that repeats e.

.zip( ... ) iterates over two iterators simultaneously, yielding the elements of both as a pair.

Macros are a bit longer to explain, so it's best to read the macro chapter in the book

like image 107
A.B. Avatar answered Feb 14 '23 08:02

A.B.