[Solved] Unified iteration over enum of vectors

#1

Hi,

I have an enum of two variants, each holding a vector of items implementing a common trait S. I’d like to write a method that returns an iterator on this vector, no matter which variant the enum is.

The code is (I put some ??? where I don’t know what to write):

trait S {
    fn x(&self);
}

struct S1 {}
impl S for S1{
    fn x(&self) {}
}

struct S2 {}
impl S for S2{
    fn x(&self) {}
}

enum Enum {
    One(Vec<S1>),
    Other(Vec<S2>),
}

fn iterate_on_enum(e: &Enum) -> ??? {
    match e {
        &Enum::One(ref s1s) => s1s.iter().map(|x| x as &S),
        &Enum::Other(ref s2s) => s2s.iter().map(|x| x as &S),
    }
}

fn main() {
}

It seems to me that I shall be able to return an iterator return trait objects, but I didn’t succeed so far in getting something to compile… Any help is appreciated :slight_smile:

#2

This works:

fn iterate_on_enum<'a>(e: &'a Enum) -> Box<Iterator<Item = &'a S> + 'a> {
    match *e {
        Enum::One(ref s1s) => Box::new(s1s.into_iter().map(|x| x as &S)),
        Enum::Other(ref s2s) => Box::new(s2s.into_iter().map(|x| x as &S)),
    }
}

The general gist is that since you’re potentially returning differently sized structs, you need to Box them up (heap allocation) in order to be able to return them. This is called a ‘trait object’.

Another thing you could do is this, although it’s slightly different conceptually:

fn iterate_on_enum<'a>(e: &'a Enum) -> IntoIter<&'a S> {
    match *e {
        Enum::One(ref s1s) => s1s.into_iter().map(|x| x as &S).collect::<Vec<_>>().into_iter(),
        Enum::Other(ref s2s) => s2s.into_iter().map(|x| x as &S).collect::<Vec<_>>().into_iter(),
    }
}

The reason you don’t need to box anything up in this scenario is that we do some ‘processing’ (via collect, which loads the references into a vector) to ensure we’re dealing with the same struct.

#3

To build on what @burns47 said, you probably want to implement IntoIterator:

impl<'a> IntoIterator for &'a Enum {
    type Item = &'a S;
    type IntoIter = Box<Iterator<Item = &'a S> + 'a>;

    fn into_iter(self) -> Self::IntoIter {
        match *self {
            Enum::One(ref s) => Box::new(s.into_iter().map(|x| x as &S)),
            Enum::Other(ref s) => Box::new(s.into_iter().map(|x| x as &S)),
        }
    }
}
#4

That’s not quite true, though. You don’t need to box up the iterator, but you’re creating a whole new vector on the heap.

#5

@birkenfeld True. I guess it’d be better phrased as “you are now able to return a statically-sized type”. Can anyone do this without a heap allocation? My intuition is telling me we should be able to do this, but I can’t articulate why…

EDIT: I suppose we could use an enum in combination with the original method above, but that feels hacky…lol

#6

Yes: You can at least always write your own Iter struct that retrieves the elements index by index. There might be a more clever solution.

#7

Either implements Iterator, so Either<slice::Iter<S1>, slice::Iter<S2>> should do the trick.

#8

Thanks for all the quality answers! The IntoIterator approach was indeed the good one, I guess I had somehow discarded it in my mind because I thought “IntoIterator consumes ownership” whereas it’s obviously OK to consume the ownership here while we’re dealing with references, and not the actual structs.

The Either based answer is quite elegant, but unfortunately my enum has more than two members, so it won’t work for me.