What is the idiomatic way to implement `IntoIterator` when some items need to be substituted?

Viewed 143

I have a custom collection like this:

struct VecChoice<T> {
    v1: Vec<T>,
    v2: Vec<T>,
    use_v1: Vec<bool>,
}

in the impl I can iterate this collection like this:

fn foo(&self, ...) {
    let item_refs: Vec<_> = (0..self.v1.len()).map(|i| {
        if self.use_v1[i] {
            &self.v1[i]
        } else {
            &self.v2[i]
        }
    });
    // ... do whatever I want with chosen references
}

However, I am failing to make it iterable:

impl<'a, T> IntoIterator for &'a VecChoice<T> {
    type Item = &'a T;

    // this fails because the trait `Sized` is not implemented for `(dyn FnMut(usize) -> Self::Item + 'static)`
    type IntoIter = Map<usize, dyn FnMut(usize) -> Self::Item>;

    fn into_iter(self) -> Self::IntoIter {
        (0..self.v1.len()).map(|i| {
            if self.use_v1[i] {
                &self.v1[i]
            } else {
                &self.v2[i]
            }
        })
    }
}

I could probably collect results into a Vec<&T> as above, then use its into_iter, but I suspect there should be a way to do it without constructing intermediate Vec.

3 Answers

The closure that you have passed to map actually does have a size. The problem though is that this type isn't nameable. You've tried to solve that with dyn, which isn't quite the right solution because the closure is sized but dyn makes it so that it isn't. dyn would be appropriate if there were different possible sizes, but then you'd have to put it behind a pointer of some kind so that the IntoIter type is Sized.

This is one of those cases where it is probably better to implement the Iterator manually, rather than using combinators.

struct VecChoiceIter<'a, T> {
    index: usize,
    vec_choice: &'a VecChoice<T>,
}

impl<'a, T> Iterator for VecChoiceIter<'a, T> {
    type Item = &'a T;
    fn next(&mut self) -> Option<Self::Item> {
        if self.index == self.vec_choice.v1.len() {
            None
        } else {
            let i = self.index;
            self.index += 1;
            let use_v1 = self.vec_choice.use_v1[i];
            if use_v1 {
                Some(&self.vec_choice.v1[i])
            } else {
                Some(&self.vec_choice.v2[i])
            }
        }
    }
}

This gives you a Sized and nameable type that you can use for the IntoIterator implementation:

impl<'a, T> IntoIterator for &'a VecChoice<T> {
    type Item = &'a T;
    type IntoIter = VecChoiceIter<'a, T>;
    fn into_iter(self) -> Self::IntoIter {
        VecChoiceIter { index: 0, vec_choice: self }
    }
}

There are some interesting RFCs in progress that could make this work more like how you originally wanted. In particular RFC-2515. This would let you write your IntoIterator implementation as you originally tried, but without having to name the type (playground - nightly):

impl<'a, T> IntoIterator for &'a VecChoice<T> {
    type Item = &'a T;

    // This is an "existential" type. That is, tell the compiler that there is 
    // exactly one possibility for what this type can be, which it can infer 
    // from the usage.
    type IntoIter = impl Iterator<Item = Self::Item>;

    fn into_iter(self) -> Self::IntoIter {
        (0..self.v1.len()).map(move |i| {
            if self.use_v1[i] {
                &self.v1[i]
            } else {
                &self.v2[i]
            }
        })
    }
}

It's often very tempting to try to make an iterator out of a pre-made collection, but unfortunately this tends to run into a practical problem a lot of the time: you need some way to store an offset into that collection, so you serve the right chunk of data out of it when next is called. Consequently, you almost always need to provide some custom iterator type.

In this case, you can do so like this:

struct VecChoice<T> {
    v1: Vec<T>,
    v2: Vec<T>,
    use_v1: Vec<bool>,
}

struct VecChoiceIter<'a, T> {
    off: usize,
    collection: &'a VecChoice<T>,
}

impl<'a, T> Iterator for VecChoiceIter<'a, T> {
    type Item = &'a T;

    fn next(&mut self) -> Option<Self::Item> {
        let off = self.off;
        self.off += 1;
        if *self.collection.use_v1.get(off)? {
            self.collection.v1.get(off)
        } else {
            self.collection.v2.get(off)
        }
    }
}

impl<'a, T> IntoIterator for &'a VecChoice<T> {
    type Item = &'a T;

    type IntoIter = VecChoiceIter<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        VecChoiceIter {
            off: 0,
            collection: self,
        }
    }
}

Note that in this case, I've switched use_v1 to a Vec<bool>, because this is not C and only booleans can be used in conditionals.

You could also do the conversion up front and store it in its own Vec, but in my experience people don't expect creating an iterator, whether by calling iter or into_iter, to be expensive. Iterators are pretty fundamental in Rust, and as a consequence it's very common for folks to create lots of them, often implicitly, and making those functions be expensive would be undesirable in many cases.

Probably the most simple way is to use .zip() and return an opaque impl Iterator from a method on the type (so you don't have to write out the actual type):

struct VecChoice<T> {
    v1: Vec<T>,
    v2: Vec<T>,
    use_v1: Vec<bool>,
}

impl<T> VecChoice<T> {
    fn iter(&self) -> impl Iterator<Item = &T> {
        self.v1
            .iter()
            .zip(self.v2.iter())
            .zip(self.use_v1.iter())
            .map(|((v1, v2), use_v1)| if use_v1 { v1 } else { v2 })
    }
}

This will iterate over all three Vec (actually the shortest of them) and return either from v1 or v2.

Notice that I switched use_v1 from a Vec<T> to a Vec<bool>, which seems to be what you have, given the way you use it.

Related