Iterators: shorten call to map

Given the following example...

fn main() {
    // This works:                                ~~~~~~~~~~~~~~~~~
    for (a, b, c) in (0..2).zip((0..2)).zip(0..2).map(|t| t.flat()) {
        println!("{} {} {}", a, b, c)
    }
    
    // This is what I want:                       ~~~~~~
    for (a, b, c) in (0..2).zip((0..2)).zip(0..2).flat() {
        println!("{} {} {}", a, b, c)
    }
}

trait Flatten<T> {
    fn flat(self) -> T;
}

impl<A, B, C> Flatten<(A, B, C)> for ((A, B), C) {
    fn flat(self) -> (A, B, C) {
        let ((a, b), c) = self;
        (a, b, c)
    }
}

(Playground)

...what do I have to implement to make it work?

Itertools already does this in some capacity. Specifically to support the iproduct!() macro (izip!() uses a different mechanism actually).

Itertools has the cons_tuples adaptor for the ((x, y), z) → (x, y, z) flattening (also for some higher arities). It simply supports the nesting that arisies in iproduct.

You may consider using itertools izip!() (macro) or multizip (function), both are interfaces to the same thing. However, using libstd .zip() has the benefit of using its zip specializations that improve performance and code generation for zip when slice iterators are involved.

1 Like

Avoiding the nesting with itertools is a good idea, but to answer your question directly:

The right-hand side in for is an iterator (or more precisely anything that implements IntoIterator), so it's not the right place to call anything on the elements, because in that expression the elements don't exist yet.

Your code should work if you call it like this:

for tuple in (0..2).zip((0..2)).zip(0..2) {
    let (a,b,c) = tuple.flat();
}

And to make .flat() work on an iterator you need to build your own Iterator that wraps the other iterators. You'd need two things: a trait that adds .flat() method to iterators, and your own struct that implements the Iterator trait. Look how stdlib does it, e.g. for enumerate(): mod.rs - source

2 Likes

I also recommend macro-based solutions like izip! for the specific case presented here. But more generally, if you have a map call that you'd like to factor out, you need to write a fair bit of code; 10~25 lines depending on how committed you are.

First, you need an extension trait. Basically, something like

trait<T> FlattenIterator<T>: Iterator
 where Self::Item: Flatten<T>
{
    // the easy way
    fn flat(self) -> Box<Iterator<T>>;

    // the Better way (not yet stable)
    fn flat(self) -> impl Iterator<T>;

    // the Right way (Iterator adapters)
    // (at least, Right for utility methods like `flat`)
    fn flat(self) -> Flat<Self,T>;
}

The primary advantages of the "right" way is increased ability to take advantage of optimizations in the stdlib iterators, and your Iterator will be able to provide len() and rev() where applicable.

What does Flat look like? The prototypical example of a custom map adapter is std::iter::Cloned. Take a look at how the std library's iter.rs used to look, back when it was still small enough to fit in a single module. Search for Cloned and delete everything else. You will see implementations for Iterator, ExactSizeIterator, and DoubleEndedIterator. (modern rust also has FusedIterator). You can pretty much copy this with very little modification. (Notice your type will also need a PhantomData member for T)

1 Like

Thx for all your replies!!

Finally I decided to make an own implementation on the basis of multizip(...) in the itertools crate from @bluss. I gave it another name (group(...)) because it follows a slightly different use case. For anyone interested in the writing process see this thread.

Here is the resulting code:

https://github.com/Bendrien/ommap/blob/1d9a958077f9b1c3fb035bc2fc6099432d722633/src/iter.rs#L68-L142