CODEWITHSUNDEEP
rust
Serdnad
Serdnad

Reputation: 642

Reuse iter variable for conditional Skip, Filter, etc

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?

Upvotes: 8

Views: 1454

Answers (2)

Jmb
Jmb

Reputation: 23434

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

Upvotes: 6

user2722968
user2722968

Reputation: 16535

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()
}

Upvotes: 8

Related Questions