How can I create an iterator of &T from either a &Vec<T> or Vec<&T>?

Question

I have an enum with two variants. Either it contains a reference to a Vec of Strings or it contains a Vec of references to Strings:

enum Foo<'a> {
    Owned(&'a Vec<String>),
    Refs(Vec<&'a String>),
}

I want to iterate over references to the Strings in this enum.

I tried to implement a method on Foo, but don't know how to make it return the right iterator:

impl<'a> Foo<'a> {
    fn get_items(&self) -> Iter<'a, String> {
        match self {
            Foo::Owned(v) => v.into_iter(),
            Foo::Refs(v) => /* what to put here? */,
        }
    }
}

fn main() {
    let test: Vec<String> = vec!["a".to_owned(), "b".to_owned()];
    let foo = Foo::Owned(&test);

    for item in foo.get_items() {
        // item should be of type &String here
        println!("{:?}", item);
    }
}

playground

What is an idiomatic method to achieve this abstraction over &Vec<T> and Vec<&T>? get_items may also return something different, as long as it implements the IntoIterator trait so that I can use it in the for loop.

Answer 1

You can't just use the std::slice::Iter type for this.

If you don't want to copy the strings or vector, you'll have to implement your own iterator, for example:

struct FooIter<'a, 'b> {
    idx: usize,
    foo: &'b Foo<'a>,
}

impl<'a, 'b> Iterator for FooIter<'a, 'b> {
    type Item = &'a String;
    fn next(&mut self) -> Option<Self::Item> {
        self.idx += 1;
        match self.foo {
            Foo::Owned(v) => v.get(self.idx - 1),
            Foo::Refs(v) => v.get(self.idx - 1).map(|s| *s),
        }
    }
}

impl<'a, 'b> Foo<'a> {
    fn get_items(&'b self) -> FooIter<'a, 'b> {
        FooIter { idx: 0, foo: self }
    }
}

fn main() {
    let test: Vec<String> = vec!["a".to_owned(), "b".to_owned()];
    let foo = Foo::Owned(&test);
    for item in foo.get_items() {
        println!("{:?}", item);
    }
    let a = "a".to_string();
    let b = "b".to_string();
    let test: Vec<&String> = vec![&a, &b];
    let foo = Foo::Refs(test);
    for item in foo.get_items() {
        println!("{:?}", item);
    }
}

Answer 2

There is a handy crate, auto_enums, which can generate a type for you so a function can have multiple return types, as long as they implement the same trait. It's similar to the code in Denys Séguret's answer except it's all done for you by the auto_enum macro:

use auto_enums::auto_enum;

impl<'a> Foo<'a> {
    #[auto_enum(Iterator)]
    fn get_items(&self) -> impl Iterator<Item = &String> {
        match self {
            Foo::Owned(v) => v.iter(),
            Foo::Refs(v) => v.iter().copied(),
        }
    }
}

Add the dependency by adding this in your Cargo.toml:

[dependencies]
auto_enums = "0.6.3"