CODEWITHSUNDEEP
rustiterator
user1002430
user1002430

Reputation:

How to have a function return one of two iterators with same item type?

The following does not compile because the two iterators have different signatures. But, both iterators return the same item type and conceptually behave the same.

fn f(flag: bool) -> impl Iterator<Item = u8> {
    if flag {
        (0..8).into_iter()
    } else {
        (0..8).into_iter().map(|x| x + 1)
    }
}

I want to write a function that returns items generated from different variants of an algorithm, selected by values I pass into the function. How can I achieve this?

Ultimately, I need to use this in a no_std environment. Is this possible?

Upvotes: 5

Views: 827

Answers (1)

Michael Anderson
Michael Anderson

Reputation: 73470

I don't know of any built-in way of doing this, but implementing something that works is not overly tricky.

What you need is a type that can represent both the result of (0..8).into_iter() - which is one kind of Iterator<Item=u8> and also (0..8).into_iter().map(|x| x+1). The typical choice for this if you don't want to use Boxed traits, is to create an enum with two options.

Then if you implement Iterator<Item=u8> on your enum and you're done.

Such an enum might look like:

pub enum EitherIter<AIterType, BIterType> {
    A(AIterType),
    B(BIterType),
}

impl<AIterType, BIterType> Iterator for EitherIter<AIterType, BIterType>
where
    AIterType: Iterator,
    BIterType: Iterator<Item = AIterType::Item>,
{
    type Item = AIterType::Item;
    fn next(&mut self) -> Option<<Self as Iterator>::Item> {
        match self {
            EitherIter::A(it) => it.next(),
            EitherIter::B(it) => it.next(),
        }
    }
}

Then using it you wrap each return type in one of the types.

pub fn f(flag: bool) -> impl Iterator<Item = u8> {
    if flag {
        EitherIter::A((0..8).into_iter())
    } else {
        EitherIter::B((0..8).into_iter().map(|x| x + 1))
    }
}

You can try this out in the playground.

Upvotes: 1

Related Questions