Map an iterator to multiple items

I'm looking for an idiomatic way for the .map() method of an iterator to decide how many items it wants to return. For example:

use std::iter::FromIterator;

fn main() {

    let v = vec![1, 2, 3, 4, 5];

    let s = Vec::<u8>::from_iter(v.into_iter().map(|e| {
        if e == 3 {
            e // I want to return e twice here
        } else {
            e
        }
    }));
    
    dbg!(s);
}

How can I change this so that the result is:

s = {
    1,
    2,
    3,
    3,
    4,
    5,
}

Use flat_map instead of map and return an iterator for an appropriate number of items, e.g. std::iter::repeat(e).take(n)

5 Likes

Ah, flat_map() is what I've been looking for.

Is there also something like repeat() where I can put in two different values or should I make a Vec?

once(1).chain(once(2)) would work. Alternatively, use arrayvec.into_iter().

Built-in arrays don't have into_iter, so you're getting slice's fallback.

2 Likes

Hm, I don't have a working solution yet. If the closure returns an iterator, I need to return an iterator in the else case as well. Also, I'd like to be able to return two different values. So I tried:

use std::collections::BTreeSet;
use std::iter::FromIterator;

fn main() {

    let v = vec![1, 2, 3, 4, 5];

    let s = Vec::<u8>::from_iter(v.into_iter().flat_map(|e| {
        if e == 3 {
            [1u8,2u8].into_iter()
        } else {
            [e].into_iter()
        }
    }));
    
    dbg!(s);
}

but that doesn't work. Apparently array.into_iter() is an iterator of references? :confused:

Ok, after re-reading your answer I tried:

use std::iter::once;
use std::collections::BTreeSet;
use std::iter::FromIterator;

fn main() {

    let v = vec![1, 2, 3, 4, 5];

    let s = Vec::<u8>::from_iter(v.into_iter().flat_map(|e| {
        if e == 3 {
            once(1).chain(once(2))
        } else {
            once(e)
        }
    }));
    
    dbg!(s);
}

but:

error[E0308]: `if` and `else` have incompatible types
  --> src/main.rs:13:13
   |
10 | /         if e == 3 {
11 | |             once(1).chain(once(2))
   | |             ---------------------- expected because of this
12 | |         } else {
13 | |             once(e)
   | |             ^^^^^^^ expected struct `std::iter::Chain`, found struct `std::iter::Once`
14 | |         }
   | |_________- `if` and `else` have incompatible types
   |
   = note: expected type `std::iter::Chain<std::iter::Once<{integer}>, std::iter::Once<{integer}>>`
            found struct `std::iter::Once<{integer}>`

I forgot about the Vec solution, that works by the way (not sure if it's good):

use std::iter::FromIterator;

fn main() {

    let v = vec![1, 2, 3, 4, 5];

    let s = Vec::<u8>::from_iter(v.into_iter().flat_map(|e| {
        if e == 3 {
            vec![11,12].into_iter()
        } else {
            vec![e].into_iter()
        }
    }));
    
    dbg!(s);
}

If the max number of items is very small like in your example, you can chain Option<T> values like this:

if e == 3 {
    once(1).chain(Some(2))
} else {
    once(e).chain(None)
}

Playground

Or if you like the flexibility of the vec! approach but need to reduce the number of allocations, you can use a crate like arrayvec or tinyvec.

2 Likes

Personally I would suggest just using Vec:

you already figured this out
fn main() {
    let v = vec![1, 2, 3, 4, 5];

    let s: Vec<u8> = v
        .into_iter()
        .flat_map(|e| {
            if e == 3 {
                vec![e, e] // I want to return e twice here
            } else {
                vec![e]
            }
        })
        .collect();

    dbg!(s);
}

(note that FromIterator is usually used indirectly through Iterator::collect, not called explicitly)

But you could use the either crate instead, to get a type that implements Iterator one of two ways:

use either::Either;

fn main() {
    let v = vec![1, 2, 3, 4, 5];

    let s: Vec<u8> = v
        .into_iter()
        .flat_map(|e| {
            if e == 3 {
                Either::Left(std::iter::repeat(e).take(2))
            } else {
                Either::Right(std::iter::once(e))
            }
        })
        .collect();

    dbg!(s);
}
2 Likes

Sort of -- since arrays don't have any into_iter() method, this call gets coerced to a slice, and then you're calling the slice implementation of IntoIterator with type Item = &'a T.

This is why we have a comparability hazard with adding IntoIterator for arrays. Code can and does exist today that expects stuff like [1, 2, 3].into_iter() to become slice::Iter, and would break if that suddenly resolved to array::IntoIter instead. There's a compiler warning about this, but still no agreement about how/when we can move forward.

2 Likes