Efficient Rust collector of Results holding Vecs - rust

Question

I'm learning Rust, and I've come upon the following pattern which collapses an iterator of Result<Vec<_>, _>s to a single big Vec<_>, failing if any of the results from the iterator failed:
fn accumulate<T, E>(it: impl Iterator<Item = Result<Vec<T>, E>>) -> Result<Vec<T>, E> {
let mut result = Vec::new();
for mut ts in it {
result.append(&mut ts?)
}
Ok(result)
}
I assume a very short "functional-style" version of this function can be written, and I'm struggling to find it. Morally, I'd like to do something like
it.map(|v| v?.into_iter()).flatten().collect()
but this doesn't typecheck. By running small examples, I think the point of the flatten there is to silently drop error results, but I'd instead like to somehow "map the flatten under the Results". I know also that in general you couldn't collect, say, an iterator of type
impl Iterator<Item = Result<impl Iterator<Item = T>, Error>>
into an iterator
Result<impl Iterator<Item = impl Iterator<Item = T>>, Error>
since you need to have actually done all of the computations in the outer iterator in order to know the final result. Nonetheless, it seems that you can make this work in this special case, when you want to .flatten() and then .collect() right after.
Finally, I can see that that collect() gives me a way to build a vector of vectors from it, and then I could flatten this vector into the single big vector I want. But this has many needless memory allocations.
Can the standard library help you do this in an efficient, Rust-ic way?

Answer 1

I think I would start with try_fold, as it can deal with Result and stop on Err:
fn acc2<T, E>(mut it: impl Iterator<Item = Result<Vec<T>, E>>) -> Result<Vec<T>, E> {
it.try_fold(
Vec::new(),
|mut vec, res_ts: Result<Vec<_>, E>| {
res_ts.map(move |mut ts| { // map preserves Err
// In case of Ok(...), append to already found elements
vec.append(&mut ts);
vec
})
}
)
}
https://play.rust-lang.org/?version=stable&mode=debug&edition=2018&gist=f6f738ddedecda1875df283f221dbfdc
It turns out, Itertools already has fold_results that should do what you want:
fn acc3<T, E>(mut it: impl Iterator<Item = Result<Vec<T>, E>>) -> Result<Vec<T>, E> {
it.fold_results(
Vec::new(),
|mut vec, mut ts| {
vec.append(&mut ts);
vec
}
)
}

Answer 2

To achieve this only using iterator methods:
use std::iter::{self, Iterator};
pub fn accumulate<T, E>(it: impl Iterator<Item = Result<Vec<T>, E>>) -> Result<Vec<T>, E> {
it.flat_map(|v| {
v.map_or_else(
|e| Iter::A(iter::once(Err(e))),
|t| Iter::B(t.into_iter().map(Ok)),
)
})
.collect()
}
// Utility enum that can be generated by the #[auto_enum] derive macro
enum Iter<T, A: Iterator<Item = T>, B: Iterator<Item = T>> {
A(A),
B(B),
}
impl<T, A: Iterator<Item = T>, B: Iterator<Item = T>> Iterator for Iter<T, A, B> {
type Item = T;
fn next(&mut self) -> Option<T> {
match self {
Self::A(a) => a.next(),
Self::B(b) => b.next(),
}
}
}
This uses flat_map to yield either an iterator of Oks or an iterator of an Err for each entry.
This is semantically equivalent to your control flow code using for loop.
Playground: https://play.rust-lang.org/?version=stable&mode=debug&edition=2018&gist=68558e27900940476e443d670a120e91
See auto_enums for deriving an enum delegating Iterator variants.
Alternatively, you can use either::Either in place of Iter, which has the same implementation for two items:
https://docs.rs/either/1.5.3/either/enum.Either.html#impl-Iterator