如何在Rust中实现函数的向量化（元素级操作转换）

问题 #

我想把一个操作两个元素的函数，转换成对两个集合进行元素级操作的函数。比如把两数相加的函数，改成接收两个数值集合、逐元素相加的函数。为了避免代码重复，我用闭包实现了vectorize函数：

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()} }

但编译器报错：

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

问题出在map方法会把闭包f从返回的闭包里移出，但f是FnMut类型，本来应该能多次用于不同迭代器。请问该怎么修改这个函数解决编译错误？

解决方案 #

核心问题是返回的闭包需要多次复用f，但当前写法每次调用闭包都会把f移动到map里，导致后续无法再使用。可以通过以下两种方式解决：

方式一：用RefCell实现内部可变性 #

把f包裹在RefCell中，让闭包可以通过可变引用多次访问它，同时保持返回类型为Fn：

use std::cell::RefCell;

fn vectorize<F, H, U, J>(f: F) -> impl Fn(H, H) -> J
where
    F: FnMut((H::Item, H::Item)) -> U + 'static,
    H: IntoIterator,
    J: FromIterator<U>,
{
    let f = RefCell::new(f);
    move |x, y| {
        x.into_iter()
            .zip(y)
            .map(|args| f.borrow_mut()(args))
            .collect()
    }
}

• RefCell提供内部可变性，允许在共享引用下获取可变引用，适配Fn闭包只能捕获不可变环境的要求。
• move关键字把RefCell<F>转移到返回的闭包中，确保闭包拥有f的所有权。
• 每次调用map时，通过borrow_mut()获取f的可变引用，调用后自动释放，保证可以多次复用。

方式二：将返回类型改为FnMut #

如果不需要返回Fn类型（允许闭包修改自身环境），可以直接把返回类型改为impl FnMut(H, H) -> J，同时调整捕获方式：

fn vectorize<F, H, U, J>(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(|args| f(args))
            .collect()
    }
}

• FnMut闭包允许捕获并修改环境中的可变变量，这里f会以可变引用的方式捕获（而非移动），每次调用闭包时都能复用f。
• 这种方式更轻量，不需要额外的RefCell开销，但返回的闭包不能被多个地方同时调用（因为FnMut不支持共享调用）。

测试示例 #

用加法函数测试上面的实现：

fn add((a, b): (i32, i32)) -> i32 {
    a + b
}

fn main() {
    // 测试方式一
    let vec_add = vectorize(add);
    let result1 = vec_add(vec![1,2,3], vec![4,5,6]);
    let result2 = vec_add(vec![7,8], vec![9,10]);
    assert_eq!(result1, vec![5,7,9]);
    assert_eq!(result2, vec![16,18]);

    // 测试方式二
    let mut vec_add_mut = vectorize(add);
    let result3 = vec_add_mut(vec![1,2], vec![3,4]);
    assert_eq!(result3, vec![4,6]);
}

内容的提问来源于stack exchange，提问作者YEp d