函数式编程: 迭代器和闭包

函数式编程在许多语言中都有, 因此 Rust 也借鉴了许多函数式编程的相关技术. 函数式编程中通常都将函数作为参数/返回值, 或将函数作为变量在之后执行等.

之前讲的模式匹配/Enum 等都深受函数式编程风格的影响, 且掌握闭包和迭代器是编写成熟/快速的 Rust 代码的必要基础.

闭包: 能够捕获上下文的匿名函数

Rust 中闭包是能够作为函数参数/返回值或存放到变量的匿名函数.

开发者可以在一个位置创建闭包, 并在另外一个位置的上下文环境中调用. 和函数不同的是, 闭包可以捕捉它定义位置的上下文环境中的值.

一个简单的例子

#[derive(Debug, PartialEq, Copy, Clone)]
enum ShirtColor {
    Red,
    Blue,
}

struct Inventory {
    shirts: Vec<ShirtColor>,
}

impl Inventory {
    fn giveaway(&self, user_preference: Option<ShirtColor>) -> ShirtColor {
        user_preference.unwrap_or_else(|| self.most_stocked())
    }

    fn most_stocked(&self) -> ShirtColor {
        let mut num_red = 0;
        let mut num_blue = 0;

        for color in &self.shirts {
            match color {
                ShirtColor::Red => num_red += 1,
                ShirtColor::Blue => num_blue += 1,
            }
        }
        if num_red > num_blue {
            ShirtColor::Red
        } else {
            ShirtColor::Blue
        }
    }
}

fn main() {
    let store = Inventory {
        shirts: vec![ShirtColor::Blue, ShirtColor::Red, ShirtColor::Blue],
    };

    let user_pref1 = Some(ShirtColor::Red);
    let giveaway1 = store.giveaway(user_pref1);
    println!(
        "The user with preference {:?} gets {:?}",
        user_pref1, giveaway1
    );

    let user_pref2 = None;
    let giveaway2 = store.giveaway(user_pref2);
    println!(
        "The user with preference {:?} gets {:?}",
        user_pref2, giveaway2
    );
}

注意到上面 giveaway 函数中 unwrap_or_else 方法就用到了闭包, unwrap_or_else 是标准库中提供的 Option<T> 对象方法. 该闭包没有参数, 有返回值, 且在闭包中使用了 self.

闭包类型推断和注解

和函数不同, 闭包一般都不需要显式指定参数/返回值类型, 因为编译器会进行自动类型推断.

不过如果将闭包赋值给变量, 则我们可以显式写类型, 这样可以增加代码可读性:

let expensive_closure = |num: u32| -> u32 {
    println!("calculating slowly...");
    thread::sleep(Duration::from_secs(2));
    num
};

闭包定义的语法可以选择如下几个(添加空白以便于阅读):

// 函数定义
fn  add_one_v1   (x: u32) -> u32 { x + 1 }
// 闭包定义: 完整标记类型
let add_one_v2 = |x: u32| -> u32 { x + 1 };
// 闭包定义: 类型自动推断
let add_one_v3 = |x|             { x + 1 };
// 如果闭包中仅单行表达式, 可以省略大括号
let add_one_v4 = |x|               x + 1  ;

注意在编译器类型推断时以第一次调用闭包时为准, 如果闭包类型推断有歧义, 比如两次调用使用不同参数类型, 则会编译报错, 如下所示:

fn main() {
    let example_closure = |x| x;

    // 类型推断: 闭包参数为 String 类型
    let s = example_closure(String::from("hello"));
    
    // 编译报错
    let n = example_closure(5);
}

闭包的上下文值捕捉: 捕捉引用或传递所有权

闭包可以在所处上下文中捕捉值, 有三种方式:

• 不可变引用(immutable reference)
• 可变引用(mutable reference)
• 传递所有权(move ownership)

一般情况下, 闭包会自动根据闭包体中如何使用上下文中的值决定使用三种方式中的哪种来捕捉.

如下是不可变引用的捕捉:

fn main() {
    let list = vec![1, 2, 3];
    println!("Before defining closure: {:?}", list);

    // 可以将闭包赋值给变量
    let only_borrows = || println!("From closure: {:?}", list);

    println!("Before calling closure: {:?}", list);

    // 捕捉 list 的不可变引用
    only_borrows();

    // 闭包执行后 list 仍然可用
    println!("After calling closure: {:?}", list);
}

捕捉值的可变引用:

fn main() {
    let mut list = vec![1, 2, 3];
    println!("Before defining closure: {:?}", list);

    let mut borrows_mutably = || list.push(7);

    // 1. 在这里没有像上个例子那样添加一行 println, 因为这样会出现一个 list 的不可变引用, 导致编译失败

    // 2. 捕捉 list 的可变引用
    borrows_mutably();
    
    // 3. 由于没有再次使用 borrows_mutably, list 的可变引用出作用域, 因此下面才能继续使用 list 的不可变引用

    // 4. 闭包执行后 list 仍然可用, 这里是 list 的不可变引用
    println!("After calling closure: {:?}", list);
}

还可以将变量 move 到闭包中, 这种技术在多线程环境下的闭包使用时非常常见:

use std::thread;

fn main() {
    let list = vec![1, 2, 3];
    println!("Before defining closure: {:?}", list);

    thread::spawn(move || println!("From thread: {:?}", list))
        .join()
        .unwrap();
}

上述在闭包参数列表前使用 move 关键字表示闭包体中使用的上下文值是直接 move 进闭包的. 因为新线程和主线程的执行完成先后顺序是不确定的, 如果主线程先执行完毕, 则 list 就会无效, 因此需要将 list move 到新线程中.

如果我们不使用 move, 闭包仍然捕捉的是不可变引用, 则可能出现 list 被提前释放的问题, 因此会编译报错, 报错信息如下:

error[E0373]: closure may outlive the current function, but it borrows `list`, which is owned by the current function

闭包中值的 Ownership 传出以及 Fn 族 Trait

闭包体中可以进行如下的上下文值操作:

• 将值的所有权传出或再传出
• 修改可变引用值
• 仅读取不可变引用值
• 不捕捉值

而闭包体对值的操作不同, 就对应了闭包实现不同的 Fn Trait. 而如果要在函数参数/返回值或在结构体的成员表示闭包, 则实际就是使用 Fn 族 Trait 来表示具体的闭包类型的.

根据闭包体对上下文值的操作, 闭包会自动实现如下 Fn Trait 的一个或多个(是叠加的关系):

1. FnOnce: 表示可以被调用一次的闭包类型, 所有的闭包都应至少实现了这个 Trait.

   如果是将内部的值 move 出去的闭包, 则它只会实现 FnOnce Trait. 因为它不能被重复调用.(比如一个闭包被存放在变量 c1 中且是仅 FnOnce 的, 则调用c1() 一次后就不能再调用了)

2. FnMut: 表示不会将捕捉的上下文值 move 出 body 的闭包, 但可能在 body 中修改被捕捉的值. 这种闭包可以被重复调用.

3. Fn: 表示不会将捕捉的上下文值 move 出 body 且不会在 body 中修改被捕捉值的, 或根本不捕捉值的闭包. 这种闭包可以被多次调用也不会改变其上下文值, 因此可以在多线程中并发调用这种闭包.

结合上面内容, 根据 body 中对值的处理, 闭包类型总结如下:

• 将 body 中值的所有权 move 出去: FnOnce, 无法重复调用
• 修改可变引用值: FnMut, 可以重复调用, 能够改变上下文值.
• 仅读取不可变引用值或不捕捉值: Fn, 可以多线程安全地重复调用, 不能改变上下文值.

在 Rust 标准库中, Fn 族定义:

pub trait FnOnce<Args> { ... }
pub trait FnMut<Args>: FnOnce<Args> { ... }
pub trait Fn<Args>: FnMut<Args> { ... }

比如标准库中 unwrap_or_else 的例子:

impl<T> Option<T> {
    pub fn unwrap_or_else<F>(self, f: F) -> T
    where
        F: FnOnce() -> T
    {
        match self {
            Some(x) => x,
            None => f(),
        }
    }
}

在 unwrap_or_else 方法中额外定义了 F 泛型参数且指定了 FnOnce 类型, 因为 f 闭包调用后, f 会将自己 body 中值所有权 move 出去(不考虑 Copy 类型的值还是非 Copy), 且 f 只能最多被调用一次(没有其他 Fn 协议实现).

在 Rust 中, 函数也可以实现上述三种协议, 如果我们不需要对上下文环境值进行捕捉, 则在传递闭包参数的时候可以传递函数名字代替. 比如在 Option<Vec<T>> 上可以这样调用 unwrap_or_else(Vec::new).

下面再来看标准库提供的用于切片的 sort_by_key 方法, 它的闭包类型为 FnMut, 闭包参数中获取到一个当前处理的切片内元素, 返回一个用于排序的 K 值:

#[derive(Debug)]
struct Rectangle {
    width: u32,
    height: u32,
}

fn main() {
    let mut list = [
        Rectangle { width: 10, height: 1 },
        Rectangle { width: 3, height: 5 },
        Rectangle { width: 7, height: 12 },
    ];
    // 闭包参数是 list 的切片中一个 item, 返回 item 中用于排序的 Key.
    list.sort_by_key(|r| r.width);
    println!("{:#?}", list);
}

sort_by_key 的闭包可能捕捉上下文中的值进行改变, 且会被调用多次, 因此使用 FnMut 类型的闭包. 如果使用 Fn 可能灵活性不强(比如想在闭包中记录一下排序调用了多少次这样), 因此标准库中没有使用这个作为 sort_by_key 的闭包类型.

迭代器(iterator)

迭代器用于处理一系列的元素, 负责将每个元素列出, 并判断元素序列是否结束.

在 Rust 中迭代器是 lazy 的, 如果没有在迭代器上调用处理方法来消费迭代器, 则迭代器不会进行任何操作:

let v1 = vec![1, 2, 3];
// 获取迭代器, 但不会进行实际的操作.
let v1_iter = v1.iter();

实际上如果在 Vec 上使用 for-in 循环, 内部实现就是先创建迭代器并消费创建的迭代器.

我们可以将这个过程分开:

let v1 = vec![1, 2, 3];

// 创建迭代器
let v1_iter = v1.iter();
// 消费迭代器
for val in v1_iter {
    println!("Got: {}", val);
}

Iterator Trait 和 next 方法

Rust 标准库中所有的迭代器都实现了 Iterator Trait:

pub trait Iterator {
    type Item;

    fn next(&mut self) -> Option<Self::Item>;

    // ...
}

其中 type Item 是为 Trait 指定关联数据类型, 使用 Item 代表序列的元素类型, 可以将 Iterator 协议理解为其他语言中的泛型协议.

Iterator Trait 仅要求实现者提供 next 方法的实现, 其余的方法都有提供默认实现. next 方法的作用是列举"下一个元素":

fn iterator_demonstration() {
    let v1 = vec![1, 2, 3];

    // 由于调用 next 会改变迭代器内部状态, 因此迭代器需要是 mut 的
    let mut v1_iter = v1.iter();

    assert_eq!(v1_iter.next(), Some(&1));
    assert_eq!(v1_iter.next(), Some(&2));
    assert_eq!(v1_iter.next(), Some(&3));
    assert_eq!(v1_iter.next(), None);
}

在 for-in 循环中没有将迭代器设置为 mut 的原因是 for 循环会拿到 iter 的所有权在内部自己处理.

且在上述 next 方法中拿到的是元素的不可变引用.

如果我们想创建一个可以获取 v1 所有权并提供 owned 值的迭代器, 则需要调用 into_iter 方法.

如果想得到一个提供元素可变引用的迭代器, 则需要使用 iter_mut 创建.

消费(Consume)迭代器的一些方法

消费迭代器(获得迭代器的所有权)的目的是确保一个迭代器仅用在一个地方, 比如 sum 方法会消费迭代器, 这样不会出现不同地方使用到同一个迭代器而出现错误的问题.

下面是使用 sum 方法的例子:

let v1 = vec![1, 2, 3];

let v1_iter = v1.iter();
// 在 sum 内部消费了迭代器, 即 fn sum(self) {...}
let total: i32 = v1_iter.sum();

// sum 调用后就无法再使用 v1_iter 了

assert_eq!(total, 6);

其他还有许多的方法可以消费迭代器, 详见官方文档.

通过一个迭代器生成另外一个迭代器的方法

这些方法不会消费迭代器, 而是通过一个迭代器生成另外一个. 下面是这类迭代器的一个例子, 即 map 方法:

let v1: Vec<i32> = vec![1, 2, 3];

// 通过 map 创建了一个新的迭代器
v1.iter().map(|x| x + 1);

map 接收的是 FnMut 闭包, 意味着也可以在 body 中捕捉上下文值的可变引用.

但上述代码中创建的新迭代器只是"需要进行的操作"的记录, 并没有实际执行迭代操作, 因此编译时会出现"未使用"的警告.

比如可以使用 collect 消费新的迭代器从而得到一个新的 Vec:

let v2: Vec<_> = v1.iter().map(|x| x + 1).collect();

迭代器生成方法的主要用法是进行链式调用, 从而可以在创建迭代器的过程中就将元素的各种处理操作链接起来, 只有在最后调用迭代器消费方法时, 才会真正执行这些链式操作的闭包.

由于许多迭代器生成方法都接收 FnMut 闭包, 我们可以在过程中捕获上下文值的可变引用, 比如下面的 filter 生成方法例子:

let shoes = vec![1, 1, 2, 3];
let shoe_size = 1;
let res = shoes.into_iter()
    .filter(|s| s == shoe_size).collect();
// res 结果为 [1, 1]

关于开发时循环和迭代器的选择

根据代码可读性和书写便捷, 以及实际的处理规模来选择, 推荐在可以使用迭代器时优先选择它.

且由于对集合的 for-in 循环内部也是使用迭代器的, 二者在性能上没有明显区别, 可能直接使用迭代器反而快一丢丢.