What return type should I use for an iterable expression?

user_318125 0
Viewed 196

I have a function that returns something that I feel should be iterable, but so far I have been collecting it and returning it as a vector.

pub fn grid_coords() -> Vec<(i32, i32)> {
    (0..SIZE).flat_map(|y| (0..SIZE).map(move |x| (x, y))).collect()
}

I assume that this will iterate before returning, and then the code calling this will just have to iterate again.

I tried removing the .collect() and letting the compiler guide me and ended up with something looking like FlatMap<Range<i32>, Map<Range<i32>, Fn>, Fn>, is there a less ugly looking type I could use?

rust
2 Answers

If you're OK with using a nightly compiler, you can use the unstable impl Trait syntax:

#![feature(conservative_impl_trait)]

pub fn grid_coords() -> impl Iterator<Item = (i32, i32)> {
    (0..SIZE).flat_map(|y| (0..SIZE).map(move |x| (x, y)))
}
user_2345900

In current (stable) Rust the easiest way to return a "conceptual" iterator is to return a boxed trait object implementing Iterator:

pub fn grid_coords() -> Box<Iterator<Item=(i32, i32)>> {
    let real_iter = (0..SIZE).flat_map(|y| (0..SIZE).map(move |x| (x, y)));
    Box::new(real_iter)
}

Using a box effectively erases the type of the concrete iterator object with the upside that the caller doesn't need to know about the fact that it's implemented using flat_map. If you later switch grid_coords to use a different iteration primitive, such as yield once it stabilizes, the signature the function will not change and you are guaranteed not to break its callers. You could even go one step further and completely hide the box from the caller:

pub struct GridCoordIter {
    inner: Box<Iterator<Item=(i32, i32)>>
}

impl Iterator for GridCoordIter {
    type Item = (i32, i32);
    fn next(&mut self) -> Option<(i32, i32)> {
        self.inner.next()
    }
}

pub fn grid_coords() -> GridCoordIter {
    let real_iter = (0..SIZE).flat_map(|y| (0..SIZE).map(move |x| (x, y)));
    GridCoordIter { inner: Box::new(real_iter) }
}

The downside of the boxed return is that it requires a small heap allocation for each call to grid_coords, and that each call to next() goes through a vtable-style indirection, neither of which is optimized away by current Rust. Whether this is a problem in practice depends on your usage, but if grid_coords() is a very basic function that might be called millions of times, it is not to be dismissed.

It is possible to eliminate the allocation and the indirection by using a different implementation of GridCoordIter. Unfortunately it is not possible to make the inner field a FlatMap value and retain the current implementation of grid_coords. The inner FlatMap would need to refer to the function's type, and the types of closures used by grid_coords are anonymous. This can be worked around by rewriting closures as callable structs, but at that point the convenience of using flat_map and map in the first place is lost, and it is easier to just implement GridCoordIter::next with the required logic:

pub struct GridCoordIter {
    i: i32,
    j: i32,
}

impl Iterator for GridCoordIter {
    type Item = (i32, i32);

    fn next(&mut self) -> Option<(i32, i32)> {
        if self.j == SIZE {
            return None;
        }
        let coord = (self.i, self.j);
        self.i += 1;
        if self.i == SIZE {
            self.i = 0;
            self.j += 1;
        }
        Some(coord)
    }
}

pub fn grid_coords() -> GridCoordIter {
    GridCoordIter { i: 0, j: 0 }
}

Achieving both efficiency and brevity requires the impl Trait feature from unstable Rust, as shown in Francis's answer. It allows the function to directly return a value of an anonymous type that the caller doesn't need to know anything about, except that it implements a particular trait.

user_16008980
Related