Elegantly flatten Vec<Result<Vec<T>, E>> into Result<Vec<T>, E>

I love how elegantly Rust can convert from Vec<Result<.. into Result<Vec<.., but for more complex nested types I'm quickly finding myself in stinkier and stinkier code. Is this as clean as it gets?:

pub type Error = Box<dyn std::error::Error>;

fn main() {
    let a: Vec<i32> = vec![];
    let b: Vec<Result<Vec<i32>, Error>> = a
        .into_iter()
        .map(|_| Ok(vec![]))
        .collect::<Vec<_>>();
        
    let c: Result<Vec<Vec<i32>>, Error> = b.into_iter().collect();

    let d: Result<Vec<i32>, Error> = c.map(|v| v.into_iter().flatten().collect::<Vec<_>>());
}

I'd say that looks pretty clean to me.

Needing to iterate over the outer Vec twice is a bit annoying though, so if I knew there would be a lot of them I would probably write the thing as a loop in a helper function and use ? to avoid passing around the Results.

fn unroll(items: Vec<Result<Vec<i32>, Error>>) -> Result<Vec<i32>, Error> {
  let mut unrolled = Vec::new();

  for result in items {
    unrolled.extend(result?);
  }

  Ok(unrolled)
}

There should also be a specialisation for Extend with a Vec<i32> that does a simple memcpy(), so that should also be pretty fast.

4 Likes

You can also write that with try_fold:

let c: Result<Vec<i32>, Error> = b.into_iter()
    .try_fold(vec![], |mut unrolled, result| {
        unrolled.extend(result?);
        Ok(unrolled)
    });
5 Likes

If you specifically need to iterate over the success values, itertools's process_results might be useful

2 Likes

I suspect that most of the time will be taken up by reallocation, so you can get a performance improvement by using Vec::with_capacity(items.len()).

A maximally performant approach would probably reuse the allocation for items, but I don't know of any efficient way of doing that without unsafe code.

To spell out the usage of process_results:

fn unroll(items: Vec<Result<Vec<i32>, Error>>) -> Result<Vec<i32>, Error> {
    items.into_iter().process_results(|i| i.flatten().collect())
}

Rust Playground

There's also a function-style one that includes the into_iter step:

fn unroll(items: Vec<Result<Vec<i32>, Error>>) -> Result<Vec<i32>, Error> {
    process_results(items, |i| i.flatten().collect())
}

Rust Playground

4 Likes

This is beautifully simple. I'll keep coming back to this answer for years to come as a great example for how my initial functional approach is just more complexity than it's probably worth. Mutability has handsome benefits clearly shown here.

2 Likes

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.