Search code examples
functionrustiterator

Returning impl Iterator with multiple types of iterator


I've tried to distill my problem into a minimal example. In the function below, if I just have either branch of the if statement, the program compiles just fine.

fn foo(bar: bool) -> impl Iterator<Item = u32> {
    if bar {
        vec![].into_iter()
    } else {
        vec![].into_iter().map(|o| o)
    }
}

However, having both branches together as I've written above gives the following error:

error[E0308]: `if` and `else` have incompatible types
 --> src/main.rs:5:9
  |
2 | /     if bar {
3 | |         vec![].into_iter()
  | |         ------------------ expected because of this
4 | |     } else {
5 | |         vec![].into_iter().map(|o| o)
  | |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected struct `std::vec::IntoIter`, found struct `Map`
6 | |     }
  | |_____- `if` and `else` have incompatible types
  |
  = note: expected struct `std::vec::IntoIter<_>`
             found struct `Map<std::vec::IntoIter<_>, [closure@src/main.rs:5:32: 5:37]>`

For more information about this error, try `rustc --explain E0308`.

By my understanding, the problem is that even though the declared return type is impl Iterator, the compiler must pick a concrete type to use. While both branches of the if statement produce a value which is of type impl Iterator, they are different concrete types. However, I don't know what concrete type I can use to resolve this problem.


Solution

  • You are correct. impl isn't a magic catch-all solution that can hide an implementation permanently. It's just convenient syntax so that you, the programmer, don't have to think about the return type. But there still must exist one unique return type.

    Now, you could (I don't recommend this) write an enum which can be either of the two iterator types.

    // Pseudocode, I'm not looking up the actual types because they're complicated
    enum MyIterator {
      IsVectorIterator(Vec::Iter<u32>),
      IsMapIterator(iter::Map<Vec::Iter<u32>, u32, ...>),
    }
    

    and then you could write an Iterator implementation for this type that delegates to whatever is inside. That's a lot of work and a lot of boilerplate, when trait objects will do it for us.

    I recommend that you return Box<dyn Iterator<Item=u32>> from your function.

    fn foo(bar: bool) -> Box<dyn Iterator<Item = u32>> {
      if bar {
        Box::new(vec![].into_iter())
      } else {
        Box::new(vec![].into_iter().map(|o| o))
      }
    }
    

    This is a dynamic trait object, so it's entirely acceptable that we don't know the concrete type until runtime. It does require dynamic dispatch (since the trait object has to remember how to be iterated over), but it's a very cheap cost to pay in exchange for not having to mess with the types explicitly.

    Note that, depending on who owns the input vector (in your example, it's a temporary vec![], but I imagine in your real code that it's an argument with a lifetime), your dyn may have to have a lifetime argument, like

    Box<dyn Iterator<Item = u32> + 'a>
    

    where 'a is the lifetime of the input vector.