Rust函数与闭包

2023-09-25 17:10 由金笔书生吕落第发表于 #后端开发

1. 常规函数

函数都拥有显示的类型签名，其本身也是一种类型。

1.1 函数类型

自由函数

// 自由函数
fn sum(a: i32, b: i32) -> i32 {
    a+b
}
fn main() {
    assert_eq!(3, sum(1, 2))
}

关联函数与方法

struct A(i32, i32);
impl A {
    // 关联函数
    fn sum(a: i32, b: i32) -> i32 {
        a+b
    }
    // 方法: 第一个参数是self, &self或&mut self的函数
    fn math(&self) -> i32 {
        Self::sum(self.0, self.1)
    }
}

fn main() {
    let a = A(1, 2);
    assert_eq!(3, A::sum(1, 2));
    assert_eq!(3, a.math());
}

1.2 函数项类型

struct A(i32, i32);
impl A {
    // 关联函数
    fn sum(a: i32, b: i32) -> i32 {
        a+b
    }
    // 方法: 第一个参数是self, &self或&mut self的函数
    fn math(&self) -> i32 {
        Self::sum(self.0, self.1)
    }
}

fn main() {
    let a = A(1, 2);
    let add = A::sum;  // Fn item type
    let add_math = A::math;  // Fn item type
    assert_eq!(add(1, 2), A::sum(1, 2));
    assert_eq!(add_math(&a), a.math());
}

函数项类型是一个零大小的类型，会在类型中记录函数的相关信息。
枚举类型与元组结构体类型与函数项类型一样，都是零大小类型。

enum Color {
    R(i16),
    G(i16),
    B(i16),
}
// 等价于
// fn Color::R(_1: i16) -> Color { /* ...  */}
// fn Color::G(_1: i16) -> Color { /* ...  */}
// fn Color::B(_1: i16) -> Color { /* ...  */}
fn main() {
    println!("{:?}", std::mem::size_of_val(&Color::R)); // 0
}

这段代码中Color::R是一个类型构造体，等价于一个函数项。
Rust默认为函数项实现了一些trait：Copy, Clone, Sync, Send, Fn, FnMut, FnOnce

2. 函数指针

函数存放在内存的代码区域内，它们同样有地址。可以使用函数指针来指向要调用的函数的地址，将函数指针传入函数中，就可以实现将函数本身作为函数的参数。

这样传递函数的方式在C语言中非常常见。

type RGB = (i16, i16, i16);
fn color(c: &str) -> RGB {
    (1, 1, 1)
}
// 这里的参数类型fn(&str)->RGB是函数指针类型, fn pointer type
fn show(c: fn(&str)->RGB) {
    println!("{:?}", c("black"));
}

fn main() {
    let rgb = color;  // rgb属于函数项类型
    show(rgb); // (1, 1, 1), 这里发生了函数项类型到函数指针类型的隐式转换
}

上述代码中rgb是一个函数项，属于函数项类型(Fn item type)，而show函数的参数则是一个函数指针类型(Fn pointer type)。

fn main() {
    let rgb = color;  // 函数项类型
    let c: fn(&str)->RGB = rgb;  // 隐式转换为了函数指针类型
    println!("{:?}", std::mem::size_of_val(&rgb));  // 0
    println!("{:?}", std::mem::size_of_val(&c));  // 8
}

应该尽可能使用函数项类型，这样有助于享受零大小类型的优化。

3. 闭包

闭包可以捕获环境变量，而函数则不可以。

// 以下代码在Rust中会报错
fn foo() -> fn(u32) {
    let msg: String = "hello".to_string();
    fn bar(n: u32) {
        for _i in 0..n {
            println!("{}", msg);
        }
    }
    return bar;
}

fn main() {
    let func = foo();
    func(5);
}

以上代码foo函数中定义的bar函数使用了环境变量msg，然而在内部函数中使用这个环境变量是被编译器所禁止的，编译器会编译报错。这是因为Rust定义函数的语法无法指定如何捕获环境变量，因此内部定义的函数无法在编译时判断使用的环境变量的生命周期是否合法，也因此Rust不允许在内部函数中使用环境变量。

想要实现以上功能就需要使用闭包。闭包在Rust中其实是一种语法糖，闭包的写法如下所示

fn foo() -> impl Fn(u32) -> () {
    let msg: String = "hello".to_string();
    let bar = move |n: u32| -> () {
        for _i in 0..n {
            println!("{}", msg);
        }
    };
    return bar;
}

fn main() {
    let func = foo();
    func(5);
}

bar是一个完整的闭包定义，其中move关键字表示捕获的环境变量所有权会被转义到闭包内，|n|是闭包的参数，-> ()表示返回值类型, {...}内是闭包的具体代码，Rust的闭包并不需要指定需要捕获的变量，闭包中使用到的环境变量会被自动捕获。Rust捕获环境变量默认是获取环境变量的引用，当使用了move关键字时，则强制捕获环境变量本身，这也就导致了所有权的转移。

3.1 闭包语法糖

Rust的闭包，实际上是语法糖，它本质上是一个实现了特定trait的匿名的struct，与闭包相关的trait有这三个：

Fn
FnMut
FnOnce

因此以上这种闭包代码它可以被展开为如下代码：

#![feature(unboxed_closures)]

fn foo() -> impl Fn(u32) -> () {
    let msg: String = "hello".to_string();

    struct ClosureEnvironment {
        env_var: String,
    }

    impl FnOnce<(u32, )> for ClosureEnvironment {
        type Output = ();

        extern "rust-call" fn call_once(self, args: (u32, )) -> Self::Output {}
    }

    impl FnMut<(u32, )> for ClosureEnvironment {
        extern "rust-call" fn call_mut(&mut self, args: (u32, )) -> Self::Output {}
    }

    impl Fn<(u32, )> for ClosureEnvironment {
        extern "rust-call" fn call(&self, args: (u32, )) -> Self::Output {
            let ClosureEnvironment { env_var } = self;
            for _i in 0..args.0 {
                println!("{}", env_var);
            }
        }
    }

    ClosureEnvironment { env_var: msg }
}

fn main() {
    let func = foo();
    func(5);
}

使用这个展开后的代码，就可以理解闭包前的move关键字的作用了，使用了move后，ClosureEnvironment结构体中的环境变量env_var保存的是String对象本身，而非引用，向其中传递环境变量msg，msg的所有权就被转移到了闭包内部。如果不使用move关键字，Rust的闭包默认会将引用传递到闭包的结构体内，而不是转移环境变量的所有权。

3.2 闭包的类型

闭包实现的trait可以为一下三种类型：

FnOnce类型

正如上面闭包展开代码所示，实现FnOnce trait中的call方法时，第一个参数的类型是self对象本身，这就会消耗闭包结构体，这也就是为什么这种闭包只能调用一次。
编译器把FnOnce的闭包类型看成函数指针。

FnMut类型

正如上面闭包展开代码所示，实现FnMut trait中的call方法时，第一个参数的类型是&mut self，是闭包对象的可变借用，不会消耗闭包结构体，切闭包函数可以对环境变量进行修改，可以被多次调用。

Fn类型

正如上面闭包展开代码所示，实现Fn trait中的call方法时，第一个参数的类型是& self，是闭包对象的不可变借用，不会消耗闭包结构体，闭包函数不可以对环境变量进行修改，可以被多次调用。

Fn: applies to closures that don’t move captured values out of their body and that don’t mutate captured values, as well as closures that capture nothing from their environment. These closures can be called more than once without mutating their environment, which is important in cases such as calling a closure multiple times concurrently.

FnMut: applies to closures that don’t move captured values out of their body, but that might mutate the captured values. These closures can be called more than once.

FnOnce: applies to closures that can be called once. All closures implement at least this trait, because all closures can be called. A closure that moves captured values out of its body will only implement FnOnce and none of the other Fn traits, because it can only be called once.

3.3 逃逸闭包与非逃逸闭包

如果使用闭包的作用域与定义闭包的作用域不同时，称该闭包为逃逸闭包，否则为非逃逸闭包。
通常如果一个函数返回值为闭包类型，则该闭包就为逃逸闭包。
逃逸闭包会遇到一个问题：如果闭包捕获了环境变量，闭包又离开了定义它的作用域，这时如果环境变量没有move或者copy到闭包中，则会出现闭包引用了原作用域中已回收变量的问题。
因此如果需要将闭包作为函数的返回值时，需要使用move将环境变量的所有权转移到闭包中，确保环境变量的生命周期在函数调用结束时不会结束。