CODEWITHSUNDEEP
rustiterator
user1002430
user1002430

Reputation:

How to conditionally chain iterators?

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?

Upvotes: 6

Views: 1906

Answers (3)

Anders Evensen
Anders Evensen

Reputation: 881

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.

Upvotes: 2

Peng Guanwen
Peng Guanwen

Reputation: 722

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>.

Upvotes: 10

cameron1024
cameron1024

Reputation: 10156

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]);
}

Upvotes: 6

Related Questions