Rust: How to short-circuit exit from a chain of iterator methods on first Err or None?

Viewed 1543

Consider a chain of iterator methods:

.iter().a().b().c()

where a produces values of type Option (or Result). Is there a way to have the whole chain return None (or (Err(_)) as soon as a yields a None (or Err(_))?

Detailed example

Given functions valid (identifying nonsensical input) and accept (an arbitrary selection criterion):

type T = u8;
type ERR = u8;

fn valid(x: &T) -> Result<T, ERR> {
    if *x < 10 { Ok(*x) } else { Err(*x) }
}

fn accept(x: &T) -> bool {
    if *x > 9 { panic!("{} should have been rejected by validator")  }
    *x % 2 == 0
}

I would like to write a function

fn count_accepted(data: &[T]) -> Result<usize, ERR>

which

  • Returns Err(ERR) as soon as the first invalid element is encountered in the input data

  • If all elements are valid, returns Ok(usize) containing the count of values that satisfied the accept criterion

Here is a solution that uses a loop:

fn count_loop(data: &[T]) -> Result<usize, ERR> {
    let mut count = 0;
    for item in data {
        valid(&item)?;
        if accept(&item) { count += 1 }
    }
    Ok(count)
}

which seems to work as required, as witnessed by these tests:

macro_rules! testem {
    ($count:path) => {
        #[test] fn empty()                 { assert_eq!($count(&[])           , Ok(0)) }
        #[test] fn all_ok_and_accepted()   { assert_eq!($count(&[2,6])        , Ok(2)) }
        #[test] fn all_ok_some_rejected()  { assert_eq!($count(&[2,3])        , Ok(1)) }
        #[test] fn one_invalid()           { assert_eq!($count(&[12])         , Err(12)) }
        #[test] fn stop_on_first_invalid() { assert_eq!($count(&[2,13,6,12,5]), Err(13)) }
    }
}

mod test_loop {testem!{super::count_loop}}

I would like to understand whether/how one could implement this behaviour using iterators rather than a loop.

Consider a related, but simpler problem: if any of the data are not valid, bail out immediately, otherwise collect all the data into a vector. In other words, remove the accept condition from the previous problem.

This problem has quite a satisfactory solution, because the FromIterator implementation of Result takes care of early termination:

fn related(data: &[T]) -> Result<Vec<T>, ERR> {
    data.iter()
        .map(valid)
        .collect()
}

mod test_related {
    #[test]
    fn stop_on_first_invalid() { assert_eq!(super::related(&[2,13,6,12,5]), Err(13))}
}

Here is an extension of related which passes the same tests as count_loop:

fn count_via_vec(data: &[T]) -> Result<usize, ERR> {
     Ok(data
        .iter()
        .map(valid)
        .filter(|x| x.is_err() || accept(&x.unwrap()))
        .collect::<Result<Vec<T>, ERR>>()?
        .len())
}

mod test_vvec {testem!{super::count_via_vec}}

However: this solution has a number of drawbacks with respect to count_loop:

  • The filter condition is very noisy.

  • The filtering step still needs to be performed when the first invalid item has been identified (unlike in the original loop implementation): the ? appears 2 lines later than it should ... if that were meaningful.

  • A vector is unnecessarily populated (unless Rust performs some cool optimization that I'm, as yet, unaware of), so the space complexity rises from O(1) to O(N).

The last point would normally be addressed by replacing .collect::<Result<Vec<T>, ERR>>()?.len()) with .count(), but this has the further detrimental effect of removing the recognition of invalid cases: they are simply counted as successes, as witnessed by the test failed by this implementation:

fn count_iterate(data: &[T]) -> Result<usize, ERR> {
    Ok(data
       .iter()
       .map(valid)
       .filter(|x| x.is_err() ||  accept(&x.unwrap()))
       .count())
}

mod test_iter {testem!{super::count_iterate}}

Can you suggest some mechanism for early returning in chains of iterator methods that can be used in cases such as this?

1 Answers

You are probably looking for std::iter::Sum which has an impl<T, U, E> Sum<Result<U, E>> for Result<T, E> where T: Sum<U>, while there are also impls for Sum for all basic integers.

So the following will Just Work:

fn valid(x: &u32) -> Result<u32, u32> {
    if *x < 10 { Ok(1) } else { Err(*x) }
}

fn count(x: &[u32]) -> Result<u32, u32> {
    x.iter()
     .map(valid)
     .sum()
}

fn main() {
    println!("{:?}", count(&[13,1,2,3]));
}

As the docs on Sum say, this will short-circuit the iterator if an error is encountered. This will include short-circuiting chained iterators.

Related