Is it possible to conditionally apply iterator transformations? [duplicate] - iterator

Question

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.

Answer 1

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.

Answer 2

The either crate provides the Either type. If both halves of Either are iterators, then so is the Either:
extern crate either;
use either::Either;
use std::iter;
fn main() {
let x: Option<i64> = None;
// Repeat x 5 times if present, otherwise count from 1 to 5
let iter = match x {
None => Either::Left(1..5),
Some(x) => Either::Right(iter::repeat(x).take(5)),
};
for i in iter {
println!("{}", i);
}
}
Like a previous answer, this still takes stack space for each concrete type you have. However, you don't need individual variables for each concrete value.
This type can also be returned from a function, unlike the trait object references. Compared to boxed trait objects, it will always use a fixed size on the stack, regardless of which concrete type was chosen.
You'll find this type (or semantic equivalent) in other places as well, such as futures::Either

Answer 3

Personally, rather than use Either, I often prefer to create a series of Option<Iterator> values that get chained together. Something like this:
playground
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 pick(x) {
println!("{}", i);
}
}
fn pick(opt_x: Option<i64>) -> impl Iterator<Item = i64> {
let iter_a = if let None = opt_x {
Some(1..5)
} else {
None
};
let iter_b = if let Some(x) = opt_x {
Some(iter::repeat(x).take(5))
} else {
None
};
iter_a.into_iter().flatten().chain(iter_b.into_iter().flatten())
}
It's a bit less obvious than using Either, but it avoids another crate, and sometimes it works out quite elegantly.

Answer 4

Here is a variation of #Niko's brilliant solution using a single match expression instead of several if let expressions, which may be more convenient when dealing with more conditional cases:
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 pick(x) {
println!("{}", i);
}
}
fn pick(opt_x: Option<i64>) -> impl Iterator<Item = i64> {
let mut iter_a = None;
let mut iter_b = None;
match opt_x {
None => iter_a = Some(1..5),
Some(x) => iter_b = Some(iter::repeat(x).take(5)),
}
iter_a.into_iter().flatten().chain(iter_b.into_iter().flatten())
}