Returning different types from a function

Imagine we want to write a function, which takes a vector, and returns an iterator.

fn to_iter(v: Vec<usize>) -> impl Iterator<Item=usize> {
    v.into_iter()
}

Pretty easy!

Now, let’s make this a bit more complicated. This function will take an additional boolean parameter called reverse, and if reverse is true, it will return a reverse iterator.

fn to_iter(v: Vec<usize>, reverse: bool) -> impl Iterator<Item=usize> {
    return if reverse {
        v.into_iter().rev()
    } else {
        v.into_iter()
    }
}

Oopsie! Compiler is angry. Because the return types are not the same anymore. Here is the error:

`if` and `else` have incompatible types
expected struct `Rev<std::slice::Iter<'_, _>>`
found struct `std::slice::Iter<'_, _>`

I was baffled with these error message at first. I mean, rev() also implements Iterator, right? Both arms of if else implement Iterator trait, and I’m returning impl Iterator. This should have translated into: I’m going to return something, but I’ll not give the concrete type, however, it wil implement Iterator trait.

So, what is this error message is trying to say to me?

Then I remembered, Rust compiler needs to know the size of what is being returned from functions to make optimizations. That’s why, only a single concrete type can be returned from a function.

Coming back to returning impl Iterator, it means: you can return something that implements Iterator trait, you don’t have to give me the concrete type, but whatever you return must have a single, concrete type.

And now, the compiler error made much more sense. Because Rev<std::slice::Iter<'_, _>> and std::slice::Iter<'_, _> are two different concrete types, even though both implement Iterator trait.

There are couple of options to address this. The first thing that comes to my mind is to use trait objects (Box<dyn Iterator>). However, this will do dynamic dispatching, will clutter the API, and I wanted something more simple.

Solution 1 (enum)

So, let’s wrap the iterators in an enum:

enum IterOwned {
    Reverse(std::iter::Rev<std::vec::IntoIter<usize>>),
    Normal(std::vec::IntoIter<usize>),
}

impl Iterator for IterOwned {
    type Item = usize;

    fn next(&mut self) -> Option<Self::Item> {
        match self {
            IterOwned::Reverse(iter) => iter.next(),
            IterOwned::Normal(iter) => iter.next(),
        }
    }
}

fn iterator(vec: Vec<usize>, reverse: bool) -> IterOwned {
    return if reverse {
        IterOwned::Reverse(vec.into_iter().rev())
    } else {
        IterOwned::Normal(vec.into_iter())
    };
}

fn main() {
    let my_vec = vec![1, 2, 3];
    let iter = iterator(my_vec, false);
    iter.for_each(|ele| println!("{}", ele));
}

What we did is a bit weird:

1. We already had 2 types that implemented iterator.
2. We wrapped these 2 different types into an enum (because compiler needs to know the exact size of the returned value)
3. Implemented Iterator trait for our wrapper enum

But what bothers me with this solution is, we are consuming the vectors. If the only thing that we are going to do is printing the contents, we don’t need to consume them. We can borrow these vectors.

Solution 2 (borrowing enum)

enum IterLifetime<'a> {
    Reverse(std::iter::Rev<std::slice::Iter<'a, usize>>),
    Normal(std::slice::Iter<'a, usize>),
}

impl<'a> Iterator for IterLifetime<'a> {
    type Item = &'a usize;

    fn next(&mut self) -> Option<Self::Item> {
        match self {
            IterLifetime::Reverse(iter) => iter.next(),
            IterLifetime::Normal(iter) => iter.next(),
        }
    }
}

fn iterator<'a>(vec: &'a Vec<usize>, reverse: bool) -> IterLifetime<'a> {
    return if reverse {
        IterLifetime::Reverse(vec.iter().rev())
    } else {
        IterLifetime::Normal(vec.iter())
    };
}

fn main() {
    let my_vec = vec![1, 2, 3];
    let iter = iterator(&my_vec, true);
    iter.for_each(|ele| println!("{}", ele));
}

However, this time, since we are borrowing the vectors, we had to use lifetimes.

At this point, subjective criteria comes into play. I’d argue the above implementation is good, does the job, and efficient. Yet, it is clunky and a bit verbose.

If we don’t don’t need to squeeze water from a stone, we can use trait objects here to significantly shorten the code.

Solution 3 (trait objects)

fn iterator(vec: Vec<usize>, reverse: bool) -> Box<dyn Iterator<Item = usize>> {
    return if reverse {
        Box::new(vec.into_iter().rev())
    } else {
        Box::new(vec.into_iter())
    };
}

fn main() {
    let my_vec = vec![1, 2, 3];
    let iter = iterator(my_vec, false);
    iter.for_each(|ele| println!("{}", ele));
}

See, the code is much more compact… with some costs:

1. we are making some heap allocations
2. API now has Box<dyn>, which might not be a good choice depending on your project’s needs
3. there will be a tiny runtime overhead due to dynamic dispatching

The goal of this article was to:

1. explain what exactly returning impl Trait means in Rust
2. provide different solutions to the problem at hand, and compare them

Hope you enjoyed!