How do I process single elements inside a flat_map?

Question

Is it possible to handle different numbers of elements inside a flat_map in Rust?

I have a case where based on a condition I need to either map a single enum element or a vector of enum elements.

I tried to use std::iter::once but it's not working since once and Vec::iter are different types the could not be returned from the same if condition.

enum Enum {
    Ident(String),
    Val(i32),
}

fn main() {
    let my_vec = vec![1, 2, 3];
    let condition = true; // Any condition

    my_vec
        .iter()
        .flat_map(|x| {
            if condition {
                std::iter::once(Enum::Ident("id".to_string())).into_iter() /* single Enum element */
            } else {
                vec![Enum::Val(1), Enum::Val(2)].iter() /* vector of Enum elements */
            }
        })
        .collect::<Vec<Enum>>()
}

error[E0308]: `if` and `else` have incompatible types
  --> src/main.rs:16:17
   |
13 | /             if condition {
14 | |                 std::iter::once(Enum::Ident("id".to_string())).into_iter() /* single Enum element */
   | |                 ---------------------------------------------------------- expected because of this
15 | |             } else {
16 | |                 vec![Enum::Val(1), Enum::Val(2)].iter() /* vector of Enum elements */
   | |                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected struct `std::iter::Once`, found struct `std::slice::Iter`
17 | |             }
   | |_____________- `if` and `else` have incompatible types
   |
   = note: expected type `std::iter::Once<Enum>`
            found struct `std::slice::Iter<'_, Enum>`

Sven Marnach · Accepted Answer · 2019-10-19T20:16:08.570

2

The obvious solution is to use a vector in both cases:

my_vec
    .iter()
    .flat_map(|_| {
        if condition {
            vec![Enum::Ident("id".to_string())].into_iter()
        } else {
            vec![Enum::Val(1), Enum::Val(2)].into_iter()
        }
    })
    .collect::<Vec<Enum>>();

This will cause an additional allocation for the one-element vector, but it is a simple and readable solution. If you want to avoid this overhead, one option is to use the either crate:

use either::{Left, Right};

enum Enum {
    Ident(String),
    Val(i32),
}

fn main() {
    let my_vec = vec![1, 2, 3];
    let _ = my_vec
        .iter()
        .flat_map(|&x| {
            if x > 1 {
                Left(std::iter::once(Enum::Ident("id".to_string())))
            } else {
                Right(vec![Enum::Val(1), Enum::Val(2)].into_iter())
            }
        })
        .collect::<Vec<Enum>>();
}

The crate defines a generic enum Either with the variants Left and Right. If the types of both side support iteration, the enum also supports iterations by dispatching to the respective underlying iterator types.

edited Oct 19 '19 at 20:16

answered Oct 19 '19 at 20:12

Sven Marnach

574,206
118
941
841

Is it possible to do this without allocating a vector for a single element? Maybe somehow turning it into iterator? – Ivan Oct 19 '19 at 20:15
1

Please see the latest version of the answer. – Sven Marnach Oct 19 '19 at 20:16
The `either` solution looks just like what I need. Is there a similar solution for a `match` clause (in the case of more than 2 branches) ? – Ivan Oct 19 '19 at 20:19
1

@Ivan I'm not aware of one. You can search for one yourself, but it's probably better to write your own enum in that case – it will probably make your code more readable. That said, it's unlikely the additional allocation matters – it hardly ever does, so please consider that solution as well. – Sven Marnach Oct 19 '19 at 20:22

How do I process single elements inside a flat_map?

1 Answers1