Writing a function returning an iterator on either a single value, or a list of such

Hi everyone,

I'm relatively new to Rust, and I'm currently stuck on something that does not seem complicated, but for which I haven't found any viable solution yet.

I have a type, which defines either a single element, or a range, i.e. something like:

enum ListElement {
    Single(u64),
    Range(u64, u64),
} 

and a list of such elements. I would like to create a function which expands this list and applies a function on each element of this expanded list, but using Iterators only (I don't want to create a vector since data I'm dealing with are very huge!).

Here is a code similar to my problem.

enum ListElement {
    Single(u64),
    Range(u64, u64),
} 

use ListElement::*;

fn iter_elements(elements: &[ListElement]) -> impl Iterator<Item=u64> {
    elements
        .iter()
        .map(|element| {
            match element {
                Single(val) => std::iter::once(*val*10),
                Range(start, end) => (*start..=*end).map(|i| i*10)
            }
        })
}


fn main() {

    let v = vec![Single(0), Range(2, 5), Single(10), Range(42, 45)];
    
    let expanded: Vec<u64> = iter_elements(&v).collect();
    
    assert_eq!(expanded, vec![0, 20, 30, 40, 50, 100, 420, 430, 440, 450])
}

Playground here

And of course, it does not compile, telling me:

error[E0308]: `match` arms have incompatible types
  --> src/main.rs:15:38
   |
13 | /             match element {
14 | |                 Single(val) => std::iter::once(*val*10),
   | |                                ------------------------ this is found to be of type `std::iter::Once<u64>`
15 | |                 Range(start, end) => (*start..=*end).map(|i| i*10)
   | |                                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected struct `std::iter::Once`, found struct `Map`
16 | |             }
   | |_____________- `match` arms have incompatible types
   |
   = note: expected type `std::iter::Once<u64>`
            found struct `Map<RangeInclusive<u64>, [closure@src/main.rs:15:58: 15:66]>`

And I don't understand this error. I have tried many different ways, and always the same result. To me, Map defines an Iterator, as well as iter::Once, so both should be accepted as a valid return value.
I am missing something for sure, but what?

Could you help me please?
Thanks in advance,

M.

But they are different types. You can't return two different types from a function conditionally. impl Iterator is not magic: it just means the compiler intentionally hides the concrete type from the caller of the function. It still has to codegen the function itself somehow, so it's not possible to pretend that two different types are the same.

If you want to decide between several concrete types dynamically, you can either create another enum that itself implements Iterator by delegating to its variants, or you can opt into dynamic dispatch by returning Box<dyn Iterator<Item=u64>> from the closure.

An alternative, which works only in your specific case, is to make the iterators into the same literal type, by returning a singleton range, val..=val, in the single-element branch.


There have been three other issues with your code:

  1. If you use impl Trait or Box<dyn Trait>, you have to annotate the function with the fact that the returned iterator borrows from the argument slice, by adding the same lifetime parameter to them.
  2. You have an iterator-of-iterators, so what you want to return a single iterator over integers is, if I understand correctly, flattening the structure.
  3. You can avoid always dereferencing every value by dereferencing the argument of map just once, which you can do by pattern matching. This in turn requires deriving Copy on ListElement.
2 Likes

Thanks @H2CO3 for your quick answer.

Unfortunately, I also tried during my previous investigations to use Box, but it failed.
Here is the two codes I tested:

fn iter_elements(elements: &[ListElement]) -> Box<dyn Iterator<Item=u64>> {
    Box::from(elements
        .iter()
        .flat_map(|element| {
            match element {
                Single(val) => Box::from(std::iter::once(*val*10)),
                Range(start, end) => Box::from((*start..=*end).map(|i| i*10))
            }
        })
    )
}

and

fn iter_elements(elements: &[ListElement]) -> Box<dyn Iterator<Item=u64>> {
    Box::from(elements
        .iter()
        .flat_map(|element| {
            match element {
                Single(val) => std::iter::once(*val*10),
                Range(start, end) => (*start..=*end).map(|i| i*10)
            }
        })
    )
}

But they both give me the same error.

I think I understand, but not sure how to translate it into Rust. Here is my attempts:

enum ListElementIterator {
    SingleIt(std::iter::Once<u64>),
    RangeIter(Box<dyn Iterator<Item=u64>>),
}
use ListElementIterator::*;

impl Iterator for ListElementIterator {
    type Item=u64;
    
    fn next(&mut self) -> Option<Self::Item> {
        match self {
            SingleIt(it_once) => it_once.next(),
            RangeIter(it) => it.next(),
        }
    }
}

fn iter_elements(elements: &[ListElement]) -> ListElementIterator {
    elements
        .iter()
        .map(|element| {
            match element {
                Single(val) => SingleIt(std::iter::once(*val*10)),
                Range(start, end) => RangeIter(Box::from((*start..=*end).map(|i| i*10)))
            }
        })
}

Especially, I have no clue about how to write the iter_elements function. What should it return?

Thanks

Oh, in the meantime, you kindly wrote the solutions, and shared it in the playground :grinning:

Thanks a lot!