Search code examples
rustclosures

How to "vectorize" a function in Rust


Suppose that I want to take a function that operates on two elements of a collection and turn it into a function that operates on two collections elementwise.

For example, I want to turn a function adding two numbers together into a function that takes two collections of numbers and adds the numbers up elementwise. In order to do this without code repetition, I tried to implement a function which takes in closures:

fn vectorize<F, H, U, J>(f: F) -> impl Fn(H, H) -> J
where
    F: FnMut((H::Item, H::Item)) -> U,
    H: IntoIterator,
    J: FromIterator<U>
    { |x, y| {x.into_iter().zip(y).map(f).collect()} }

The comipler gives:

error[E0507]: cannot move out of `f`, a captured variable in an `Fn` closure
  --> src/main.rs:10:40
   |
5  | fn vectorize<F, H, U, J>(f: F) -> impl Fn(H, H) -> J
   |                          - captured outer variable
...
10 |     { |x, y| {x.into_iter().zip(y).map(f).collect()} }
   |       ------                           ^ move occurs because `f` has type `F`, which does not implement the `Copy` trait
   |       |
   |       captured by this `Fn` closure

This does not work because the closure f is moved out of the returned closure by the map method. But really, such a function should be possible to write, since the f is FnMut and will be valid no matter how many times it is used on different iterators.

How can I write this type of function so that it does not cause such errors?


Solution

  • You want to change three things here.

    1. You can't convert a FnMut into a Fn so you have to return impl FnMut as well.
    2. Add move before the closure so it takes ownership of f
    3. You can't pass f to map, because that necesarily moves it, but you can't move from a closure that only implements Fn as in your code or FnMut as in the fixed code. Instead pass a reference to it (only a mutable reference also implements FnMut and to pass a mutable reference you have to make f mutable as well):
    fn vectorize<F, H, U, J>(mut f: F) -> impl FnMut(H, H) -> J
    where
        F: FnMut((H::Item, H::Item)) -> U,
        H: IntoIterator,
        J: FromIterator<U>,
    {
        move |x, y| x.into_iter().zip(y).map(&mut f).collect()
    }
    

    You might also want to add a second : IntoIterator type parameter so you can merge 2 different types of collections with possibly distinct types as well.