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

user_1002430 0
Viewed 116

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?

rust iterator
1 Answers

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.

user_2219550
Related