How to use iterators that lead to two collects

[jorgecarleitao]

Consider the loop

let r1 = vec![];
let r2 = vec![];
for i in 0..20 {
    r1.push(i*i)
    r2.push(i*i*i)
}

is there a way to write this as an iterator over i, whose collect from the two maps (|i| i*i and |i| i*i*i) returns the same semantic content as the for above and with the same performance?

[mbrubeck]

You could do something like this:

fn main() {
    let TwoVecs(r1, r2) = (0..20).map(|i| (i*i, i*i*i)).collect();
    dbg!(r1);
    dbg!(r2);
}

struct TwoVecs<T, U>(Vec<T>, Vec<U>);

impl<T, U> std::iter::FromIterator<(T, U)> for TwoVecs<T, U> {
    fn from_iter<I>(iter: I) -> Self
        where I: IntoIterator<Item = (T, U)>
    {
        let iter = iter.into_iter();
        let size = iter.size_hint().0;
        let mut r1 = Vec::with_capacity(size);
        let mut r2 = Vec::with_capacity(size);
        for (a, b) in iter {
            r1.push(a);
            r2.push(b);
        }
        Self(r1, r2)
    }
}

Playground

[Sagebati]

Without the impl FromIterator, you can also try:
let (r1,r2): (Vec<_>, Vec<_>) = 0..20.map(|i| (i*i, i*i*i)).unzip()

[mbrubeck]

Ah, of course! After all these years, I still sometimes forget some of the standard iterator methods...