How do I flatten an iterator of Result<Vec<T>,E> to return an iterator of Result<T,E> in rust?

Viewed 121

This is my test case; it doesn't compile because of incompatible match arms. How should I fix that, or what's a better way to get a flattened iterator here?

fn main() {
    let v1: Vec<Result<Vec<u32>, &str>> = vec![ Ok(vec![1,2,3]), Ok(vec![4,5,6,9]), Err("failed"), ];

    let iter1 = v1.iter();

    let flattened_iter1 = iter1.map(|x| match x {
        Ok(v) => v.iter().map(|v| Ok(v)),
        Err(e) => once(Err(e)),
    }).flatten();

    for r in flattened_iter1 {
        match r {
            Ok(i) => println!("{}", i),
            Err(i) => { println!("Error"); break },
        }
    }
}

3 Answers

One option is to use the either crate, which provides the Either<L, R> enum. An Either<L, R> can be either Either::Left(L) or Either::Right(R), and Either<L, R> implements Iterator itself when both L and R implement Iterator with the same Iterator::Item type.

This gives you similar flexibility to a boxed trait object but without the overhead of heap allocation. There is still dynamic dispatch, but handled internally by Either as a conditional instead of a function pointer call.

use std::iter::once;
use either::Either;

fn main() {
    let v1: Vec<Result<Vec<u32>, &str>> = vec![ Ok(vec![1,2,3]), Ok(vec![4,5,6,9]), Err("failed"), ];

    let iter1 = v1.iter();

    let flattened_iter1 = iter1.map(|x| match x {
        Ok(v) => Either::Left(v.iter().map(|v| Ok(v))),
        Err(e) => Either::Right(once(Err(e))),
    }).flatten();

    for r in flattened_iter1 {
        match r {
            Ok(i) => println!("{}", i),
            Err(i) => { println!("Error"); break },
        }
    }
}

(Playground)

This won't work because map and once produce two different types. What you would like would be a way of telling to Rust you don't really care that they are different types, the only thing you care is that they are both iterators. To do so, you could explicitly cast them as trait objects once(Err(e)) as dyn Iterator<Item=Result<u32, &str>>, but that wouldn't work because then Rust couldn't know their size at compile time. To make it work, you need to allocate them on the heap rather than on the stack, leading to

use std::iter::once;

fn main() {
    let v1 = vec![ Ok(vec![1,2,3]), Ok(vec![4,5,6,9]), Err("failed"), ];

    let flattened_iter1 = v1
        .iter()
        .map(|x| match x {
            Ok(v) => Box::new(v.iter().map(|v| Ok(v))) as Box<dyn Iterator<Item=_>>,
            Err(e) => Box::new(once(Err(e))) as Box<dyn Iterator<Item=_>>,
        })
        .flatten();

    for r in flattened_iter1 {
        match r {
            Ok(i) => println!("{}", i),
            Err(i) => { println!("Error"); break },
        }
    }
}

You can try this on the playground.

There are a few things to note with this approach:

  • Boxing stuff is not free, as it results in actual allocate, but it's also not super expensive. Compared to the solution where you just collect everything into a single Vec to iterate on that, this is probably cheaper;
  • Similarly, working with trait objects is not free (because they have to bring their own vtable and everything), but it's also not very expensive.
  • I didn't have to specify the actual iterator Item because Rust was able to figure it out on its own.

While in the general case I would use the either crate, in this specific case itertools has gotten you covered with flatten_ok():

use itertools::Itertools;

let flattened_iter1 = iter1.map(Result::as_ref).flatten_ok();
Related