Idiomatic way to optionally add one item to an iterator? [duplicate] - rust

Question

Let's say I have:
let it = [1, 2, 3].into_iter();
let jt = [4, 5, 6].into_iter();
let kt = [7, 8, 9].into_iter();
Then I have boolean conditions i, j and k. I want to generate an iterator that conditionally chains it, jt and kt together based on the values of i, j and k. Can I do this with just the built-in Rust Iterator functionality?

Answer 1

You can make Option into an iterator.
let it = i.then_some([1, 2, 3]).into_iter().flatten();
let jt = j.then_some([4, 5, 6]).into_iter().flatten();
let kt = k.then_some([7, 8, 9]).into_iter().flatten();
let iter = it.chain(jt).chain(kt);
If the condition is false, then condition.then_some(...) will return None, making an empty iterator. Otherwise a Some(...) is returned. into_iter().flatten() will transform Option<impl IntoIterator<Item=T>> to impl Iterator<Item=T>.

Answer 2

You're going to run into a slight issue if you want to use bare iterators:
If you write the following:
let iter = [1, 2, 3].into_iter();
let iter = if some_condition {
iter.chain([4, 5, 6])
} else {
iter
}
You'll get an error which boils down to this:
= note: expected struct `std::iter::Chain<std::array::IntoIter<_, _>, std::array::IntoIter<{integer}, 3>>`
found struct `std::array::IntoIter<_, _>`
iter has type IntoIter, but iter.chain() has type Chain<IntoIter, ...>
To get around this, you have a few options:
you can use a trait object, which behaves a bit like an interface from languages like Java, but loses some performance:
let iter = [1, 2, 3].into_iter();
let mut iter: Box<dyn Iterator<Item = i32>> = Box::new(iter);
if some_condition {
iter = Box::new(iter.chain([4, 5, 6]));
}
or, probably a better solution, if you can sacrifice laziness, is to just use a Vec:
// save heap allocations by pre-allocating the whole vec
let len = if some_condition { 6 } else { 3 };
let mut items = Vec::with_capacity(len);
items.extend([1, 2, 3]);
if some_condition {
items.extend([4, 5, 6]);
}

Answer 3

This is a good use for the either crate. Either implements Iterator when both the left and right sides also implement Iterator, so it can be easily used to chain iterators together.
Given any three iterators it, jt, and kt that iterate over the same Item, with accompanying booleans i, j, and k, you can write a function that chains them together like this:
use either::Either;
use std::iter;
fn chain<'a, I, J, K, Item>(
it: I,
jt: J,
kt: K,
i: bool,
j: bool,
k: bool,
) -> iter::Chain<
iter::Chain<Either<I, iter::Empty<Item>>, Either<J, iter::Empty<Item>>>,
Either<K, iter::Empty<Item>>,
>
where
I: Iterator<Item = Item>,
J: Iterator<Item = Item>,
K: Iterator<Item = Item>,
{
let iter = if i {
Either::Left(it)
} else {
Either::Right(iter::empty())
};
let iter = iter.chain(if j {
Either::Left(jt)
} else {
Either::Right(iter::empty())
});
let iter = iter.chain(if k {
Either::Left(kt)
} else {
Either::Right(iter::empty())
});
iter
}
Calling this function will result in an iterator conditional on the input. For example, calling
let it = [1, 2, 3].into_iter();
let jt = [4, 5, 6].into_iter();
let kt = [7, 8, 9].into_iter();
chain(it, jt, kt, true, false, true).collect::<Vec<_>>();
gives
[1, 2, 3, 7, 8, 9]
as expected.
You can try it using this playground.