Search code examples
rustiterator

What is the canonical way to take one extra element in an iterator after a predicate fails?


I have a Vec<bool> and I want to create an iterator that stops at and including the first false; How can I implement that? I've been looking at the take_while method, but it stops one element early. For example, if the input was vec![true, true, true, false, false], I want an iterator equivalent to vec![true, true, true, false].iter(), but I can't find a good way to express that in a functional way. I have considered the following methods:

// method 1
// This one triggers the side effects of do_something_with_x
// for the right elements, but it seems like bad style. It
// also doesn't give me an iterator that I can chain on, I
// would need to use the side effect creatively.
input
    .iter()
    .map(|x| {
        do_something_with_x(x);
        x
    })
    .take_while(|x| *x)
    .for_each(|_|());

// method 2
// This seems a little better, but I don't like how I have
// to iterate through it twice. My code isn't performance
// sensitive so that's not the reason, but it just looks bad to me.
let len = input.iter().position(|x| !*x);
let result = input.iter().take(len);

I have significantly simplified the problem to make it easier to understand, the vector is of a different type and I have a different predicate I'm using. I'm also thinking about a third method that would involve zipping two iterators of the same thing that are offset by one, but I'm not quite sure how to best implement that and it doesn't seem very readable regardless. What is the canonical way to solve this kind of problem?


Solution

  • You can use Iterator::scan, which produces a new iterator by calling a closure with each item and a mutable state.

    By letting the state indicate if we should end at the next item, we can stop after the first falsy item:

    input
        .iter()
        .scan(false, |is_done, item| {
            if *is_done { return None; }
            if !item { *is_done = true; }
            Some(item)
        })
        .for_each(|_| {});