Rust闭包

前面几节介绍了在Lua中定义的闭包。除此之外,Lua语言的官方实现还支持C语言闭包。我们的解释器是由Rust实现的,自然也就要改成Rust闭包。本节就来介绍Rust闭包。

Lua官方实现中的C闭包

先来看下Lua官方实现中的C闭包。C语言本身不支持闭包,所以必须依赖Lua配合才能实现闭包。具体来说,就是把Upvalue存到Lua的栈上,然后再跟C函数原型绑定起来组成C闭包。Lua通过API向C函数提供访问栈上Upvalue的方式。

下面是C闭包版本的计数器示例代码:

// 计数器函数原型
static int counter(Lua_State *L) {
    int i = lua_tointeger(L, lua_upvalueindex(1)); // 读取Upvalue计数
    lua_pushinteger(L, ++i);  // 加1,并压入栈顶
    lua_copy(L, -1, lua_upvalueindex(1));  // 用栈顶新值更新Upvalue计数
    return 1;  // 返回栈顶的计数
}

// 工厂函数,创建闭包
int new_counter(Lua_State *L) {
    lua_pushinteger(L, 0);  // 压到栈上
    
    // 创建C闭包,函数原型是counter,另外包括1个Upvalue,即上一行压入的0。
    lua_pushcclosure(L, &counter, 1);

    // 创建的C闭包压在栈顶,下面return 1代表返回栈顶这个C闭包
    return 1;
}

先看第2个函数new_counter(),也是创建闭包的工厂函数。先调用lua_pushinteger()把Upvalue计数压入到栈顶;然后调用lua_pushcclosure()创建闭包。复习一下,闭包由函数原型和Upvalue组成,这两部分分别由lua_pushcclosure()函数的后面两个参数指定。第一个参数指定函数原型counter,第二个参数1代表栈顶的1个Value是Upvalue,即刚刚压入的0。下图是调用这个函数创建C闭包前后的栈示意图:

|     |                            |         |
+-----+                            +---------+
|  i  +--\  +-C_closure------+<----+ closure |
+-----+  |  | proto: counter |     +---------+
|     |  |  | upvalues:      |     |         |
         \--+--> i           |
            +----------------+

上图最左边是把计数i=0压入到栈顶。中间是创建的C闭包,包括了函数原型和Upvalue。最右边是创建完闭包后的栈布局,闭包压入栈上。

再看上述代码中第1个函数counter(),也就是创建的闭包的函数原型。这个函数比较简单,其中最关键的是lua_upvalueindex() API,生成代表Upvalue的索引,就可以用来读写被封装在闭包中的Upvalue了。

通过上述示例中代码对相关API的调用流程,基本可以猜到C闭包的具体实现。我们的Rust闭包也可以参考这种方式。但是,Rust本身就支持闭包!所以我们可以利用这个特性更简单的实现Lua中的Rust闭包。

Rust闭包的类型定义

用Rust语言的闭包实现Lua中的“Rust闭包”类型,就是新建一个Value类型,包含Rust语言的闭包就行。

《Rust程序设计语言》中已经详细介绍了Rust的闭包,这里就不再多言。我们只需要知道Rust闭包是一种trait。具体到Lua中的Rust闭包类型就是FnMut (&mut ExeState) -> i32。然后就可以尝试定义Lua中Value的Rust闭包类型如下:

pub enum Value {
    RustFunction(fn (&mut ExeState) -> i32),   // 普通函数
    RustClosure(FnMut (&mut ExeState) -> i32), // 闭包

然而这个定义是非法的,编译器会有如下报错:

error 782| trait objects must include the `dyn` keyword

这就涉及到Rust中trait的Static Dispatch和Dynamic Dispatch了。对此《Rust程序设计语言》也有详细的介绍,这里不再多言。

然后,我们根据编译器的提示,加上dyn

pub enum Value {
    RustClosure(dyn FnMut (&mut ExeState) -> i32),

编译器仍然报错,但是换了一个错误:

error 277| the size for values of type `(dyn for<'a> FnMut(&'a mut ExeState) -> i32 + 'static)` cannot be known at compilation time

就是说trait object是个DST。这个之前在介绍字符串定义的时候介绍过,只不过当时遇到的是slice,现在是trait,这也是Rust中最主要的两个DST。对此《Rust程序设计语言》也有详细的介绍。解决方法就是在外面封装一层指针。既然Value要支持Clone,那么Box就不能用,只能用Rc。又由于是FnMut而不是Fn,在调用的时候会改变捕捉的环境,所以还需要再套一层RefCell来提供内部可变性。于是得到如下定义:

pub enum Value {
    RustClosure(Rc<RefCell<dyn FnMut (&mut ExeState) -> i32>>),

这次终于编译通过了!但是,想一想当初在介绍字符串各种定义的时候为什么没有使用Rc<str>?因为对于DST类型,需要在外面的指针或引用的地方存储实际的长度,那么指针就会变成“胖指针”,需要占用2个word。这就会进一步导致整个Value的size变大。为了避免这种情况,只能再套一层Box,让Box包含具体长度变成胖指针,从而让Rc恢复1个word。定义如下:

pub enum Value {
    RustClosure(Rc<RefCell<Box<dyn FnMut (&mut ExeState) -> i32>>>),

在定义了Rust闭包的类型后,也遇到了跟Lua闭包同样的问题:还要不要保留Rust函数的类型?是否保留区别都不大。我们这里选择了保留。

虚拟机执行

Rust闭包的虚拟机执行非常简单。因为Rust语言中闭包和函数的调用方式一样,所以Rust闭包的调用跟之前Rust函数的调用一样:

    fn do_call_function(&mut self, narg_plus: u8) -> usize {
        match self.stack[self.base - 1].clone() {
            Value::RustFunction(f) => { // Rust普通函数
                // 省略参数的准备
                f(self) as usize
            }
            Value::RustClosure(c) => { // Rust闭包
                // 省略同样的参数准备过程
                c.borrow_mut()(self) as usize
            }

测试

至此就完成了Rust闭包类型。借用了Rust语言自身的闭包后,这个实现就非常简单。并不需要像Lua官方实现那样用Lua栈来配合,也就不需要引入一些专门的API。

下面代码展示了用Rust闭包来完成本节开头的计数器例子:

fn test_new_counter(state: &mut ExeState) -> i32 {
    let mut i = 0_i32;
    let c = move |_: &mut ExeState| {
        i += 1;
        println!("counter: {i}");
        0
    };
    state.push(Value::RustClosure(Rc::new(RefCell::new(Box::new(c)))));
    1
}

相比于本节开头的C闭包,这个版本除了最后一句创建闭包的语句非常啰嗦以外,其他流程都更加清晰。后续在整理解释器API时也会优化最后这条语句。

Rust闭包的局限

上面的示例代码中可以看到,捕获的环境(或者说Upvalue)i是需要move进闭包的。这也就导致多个闭包间共享不能共享Upvalue。不过Lua官方的C闭包也不支持共享,所以并没什么问题。

另外一个需要说明的地方是,Lua官方的C闭包中是用Lua的栈来存储Upvalue,也就导致Upvalue的类型就是Lua的Value类型。而我们使用Rust语言的闭包,那Upvalue就可以是“更多”的类型,而不限于Value类型了。不过这两者之间在功能上应该是等价的:

  • Rust闭包支持的“更多”类型,在Lua中都可以用LightUserData,也就是指针来实现;虽然对于Rust来说这很不安全。
  • Lua中支持的内部类型,比如表Table,在我们的解释器中,也可以通过get()这个API获取到(而Lua的官方实现中,表这个类型是内部的,没有对外)。