Conditionally iterate over one of several possible iterators

rust

Solution

The most straightforward solution is to use a trait object:

use std::iter;

fn main() {
    let mut a;
    let mut b;

    let x: Option<i64> = None;

    // Repeat x 5 times if present, otherwise count from 1 to 5
    let iter: &mut dyn Iterator<Item = i64> = match x {
        None => {
            a = 1..5;
            &mut a
        }
        Some(x) => {
            b = iter::repeat(x).take(5);
            &mut b
        }
    };

    for i in iter {
        println!("{}", i);
    }
}

The main downside for this solution is that you have to allocate stack space for each concrete type you have. This also means variables for each type. A good thing is that only the used type needs to be initialized.

The same idea but requiring heap allocation is to use boxed trait objects:

use std::iter;

fn main() {
    let x: Option<i64> = None;

    // Repeat x 5 times if present, otherwise count from 1 to 5
    let iter: Box<dyn Iterator<Item = i64>> = match x {
        None => Box::new(1..5),
        Some(x) => Box::new(iter::repeat(x).take(5)),
    };

    for i in iter {
        println!("{}", i);
    }
}

This is mostly useful when you want to return the iterator from a function. The stack space taken is a single pointer, and only the needed heap space will be allocated.

You can also use an enum for each possible concrete iterator.

Problem

I'm trying to switch behavior based on an `Option` input to a function. The idea is to iterate based on whether or not a given `Option` is present. Here's a minimal, if silly, example: ``` use std::iter; fn main() { let x: Option<i64> = None; // Repeat x 5 times if present, otherwise count from 1 to 5 for i in match x { None => 1..5, Some(x) => iter::repeat(x).take(5), } { println!("{}", i); } } ``` I get an error: ``` error[E0308]: match arms have incompatible types --> src/main.rs:7:14 | 7 | for i in match x { | ______________^ 8 | | None => 1..5, 9 | | Some(x) => iter::repeat(x).take(5), | | ----------------------- match arm with an incompatible type 10 | | } { | |_____^ expected struct `std::ops::Range`, found struct `std::iter::Take` | = note: expected type `std::ops::Range<{integer}>` found type `std::iter::Take<std::iter::Repeat<i64>>` ``` This makes perfect sense, of course, but I'd really like to choose my iterator based on a condition, since the code in the for-loop is non-trivial and copy-pasting all of that just to change iterator selection would be pretty ugly and unmaintainable. I tried using `as Iterator<Item = i64>` on both arms, but that gives me an error about unsized types because it's a trait object. Is there an easy way to go about this? I could, of course, use `.collect()` since they return the same type and iterate over that vector. Which is a good quick fix, but for large lists seems a bit excessive.

Original source

Related problems