Reuse iter variable for conditional Skip, Filter, etc

user_2330228 0
Viewed 297

I have an iterator to which I want to conditionally apply filters, skips, etc. before ultimately calling collect().

In a different language, I might write something like

// (Rust-like pseudocode)
let mut iter = [1,2,3,4].iter();
if should_filter {
    iter = iter.filter(|x| x % 2 == 0);
}
if should_truncate {
    iter = iter.take(2);
}
iter.collect()

But since iter is of type Iter, and skip(), filter() return types Skip, Filter, I've been unable to reuse the original binding for iter. As a result, my Rust code currently looks something like this:

let iter = [1,2,3,4].iter();

// conditionally filter
if should_filter {
    let iter = iter.filter(|x| x % 2 == 0);

    // conditionally "truncate"
    if should_truncate {
        let iter = iter.take(2);
        return iter.collect();
    }

    return iter.collect();
}

// conditionally "truncate"
if should_truncate {
    let iter = iter.take(2);
    return iter.collect();
}

iter.collect()

Is there any way I can avoid this duplication?

rust
2 Answers

You can reliably trust the compiler the optimize loop invariants in this case. Since should_filter can not change while iterating, the compiler will figure out that it can check that precondition before the loop and skip the test if should_filter is true. This means you can simply put the condition into the loop - which seems inefficient - and have much cleaner code. Even if the check does not get removed from the loop body, the CPU's branch predictor will easily skip around it. Similarly you can "inline" the should_truncate condition:

fn do_stuff(
    inp: impl IntoIterator<Item = u32>,
    should_filter: bool,
    should_truncate: bool,
) -> Vec<u32> {
    inp.into_iter()
        .filter(|x| should_filter && x % 2 == 0)
        .take(if should_truncate { 2 } else { usize::MAX })
        .collect()
}
user_27229680

Easiest solution, and probably closest to what your "other language" would do under the hood, is to box the iterator:

let mut iter: Box<dyn Iterator<Item = &i32>> = Box::new ([1, 2, 3, 4].iter());
if should_filter {
    iter = Box::new (iter.filter(|x| *x % 2 == 0));
}
if should_truncate {
    iter = Box::new (iter.take(2));
}
let v: Vec<_> = iter.collect();

Playground

user_53970090
Related