Modify an Iterator Until a Condition and then Continue

Imagine you have a collection of numbers, some of which are zero. You would like to replace all the numbers up until the first zero with zeroes, and from then on make no changes to the rest of the collection.

zero_until_first_zero([1, 2, 3, 0, 1, 2, 3]) -> [0, 0, 0, 0, 1, 2, 3]
zero_until_first_zero([1, 2, 3, 4]) -> [0, 0, 0, 0]
zero_until_first_zero([0, 1, 2, 3]) -> [0, 1, 2, 3]
zero_until_first_zero([]) -> []

My friend and I came up with a couple options:

c.iter().scan(false, |found_zero, &elem| {
    *found_zero = elem == 0 || *found_zero;
    Some(if !*found_zero { 0 } else { elem })
}).collect()
c.iter().by_ref()
.take_while(|elem| *elem != &0).map(|_| &0)
.chain(c.iter().skip_while(|elem| *elem != &0))
.collect()

Also, a variation of the second version with map_while() from nightly.

Finally, a straightforward but not-as-fun version that uses a for loop.

All of these seem to lack a certain elegance. I think scan() is the right choice, but the &'s and *'s don't sit quite right for me. How would you do it in an idiomatic way?

Iterator::map can use a FnMut

Not sure what the goal is edit: looks I misread, you want idiomatic code, but in real code I'd go with something like:

for x in collection.iter_mut().take_while(|x| x != 0) {
    *x = 0;
}

I prefer to see mutations in the for body

6 Likes

Yeah, if I weren't trying to figure this out for fun and actually had to solve this problem, I'm sure I would have written that exact loop and moved on quickly :slight_smile:

This is exactly what I was looking for! I also re-wrote it with scan() using your syntax of |=, which I like better than what I originally had:

I feel compelled to note that "idiomatic" does not mean "using the maximal number of iterators". A for loop is a perfectly idiomatic and sensible solution to this problem. scan is extraneous, but passing a stateful closure to map is also kind of weird and people will look at you funny for doing it. "Idiomatic" is using the tools the language gives you in the ways they're meant to be used. for is one of those tools and this is a perfectly appropriate use for it.

3 Likes

Agreed. Idiomatic code is code that an intermediate rust user can skim over without it eliciting questions.

Maybe it wasn't clear from my answer (there's a "but" that's a leftover of the stricken text), but by "real code" I meant "idiomatic" too

Library version: https://play.rust-lang.org/?version=stable&mode=debug&edition=2018&gist=b6d69a88b22b16627c8ecf6ac3e71c83

This is a good point, and I agree. The for loop is probably the right choice, and I miswrote what I actually wanted, which was the "cool" iterator way to do it. :slight_smile:

Thanks all for the help!

1 Like