Can I specify the type for iterators based on closures?

user_40916 0
Viewed 108

I wrote a function to iterate over neighbors of cells in a 2d grid:

pub fn neighbours(
    (width, height): (usize, usize),
    (x, y): (usize, usize),
) -> impl Iterator<Item = (usize, usize)> {
    [(-1, 0), (1, 0), (0, -1), (0, 1)]
        .iter()
        .map(move |(dx, dy)| (x as i64 + dx, y as i64 + dy))
        .filter(move |&(nx, ny)| 0 <= ny && ny < height as i64 && 0 <= nx && nx < width as i64)
        .map(|(nx, ny)| (nx as usize, ny as usize))
}

Note that it returns an impl Iterator<Item = (usize, usize)>.

If I understand correctly, returning an impl would result in slower code, calling function pointers instead of compiling things down to simple loops. Right?

So wanting to specify a more exact type, I replaced the type with () to see what type the compiler infers, and it infers

std::iter::Map<std::iter::Filter<std::iter::Map<std::slice::Iter<'l, (i64, i64)>, _>, _>, _>

Where the _s stand for closures which I don't know how to specify their types.

I tried to extract the closures to structs with Fn traits but couldn't make that work and also IIUC implementing Fn traits is an unstable feature and I shouldn't use it"

rust closures types
1 Answers

If I understand correctly, returning an impl would result in slower code, calling function pointers instead of compiling things down to simple loops. Right?

Nope. Returning impl Iterator<Item = (usize, usize)> is exactly the same, codegen-wise, as returning Map<Filter<Map<...>>>.¹ The difference is twofold:

  • Changing the return type of neighbours does not require changing its signature, so impl Iterator is forward-compatible with changes to its return type.
  • impl Iterator won't unify with other opaque impl Iterator types, even if they happen to have the same underlying type. (In simple terms: the compiler won't allow you to make a Vec of impl Iterators from different sources, even if all those opaque types are the same concrete type.)

Neither of these differences has any influence on code generation or the compiler's ability to inline anything, so go ahead and use impl Iterator.

There is one case where you must still use indirect dispatch (dyn Iterator): when the function neighbours is itself part of a trait, the impl Trait syntax is not yet available (as of 1.59). The best way to solve this at the moment is to return Box<dyn Iterator>. (Note, however, that doesn't mean every call will be dynamically dispatched; the call to .next() will be, but everything "inside" that still uses easily-optimized static dispatch.)

Related questions

¹ Note that in order to actually return Map<Filter<Map<...>>>, you would still have to use impl Trait to represent the closures, since they have anonymous types.

user_36503620
Related