Rust Closure

The previous sections introduced closures defined in Lua. In addition, the official implementation of the Lua language also supports C language closures. Our interpreter is implemented by Rust, so it will naturally be changed to Rust closures. This section introduces Rust closures.

C Closures in the Official Implementation of Lua

Let's first look at the C closures in the official implementation of Lua. The C language itself does not support closures, so it must rely on the cooperation of Lua to realize closures. Specifically, the upvalue is stored on the Lua stack, and then bound to the C function prototype to form a C closure. Lua provides a way for C functions to access upvalue on the stack through API.

Here is the C closure version of the counter example code:

// counter function prototype
static int counter(Lua_State *L) {
     int i = lua_tointeger(L, lua_upvalueindex(1)); // read upvalue count
     lua_pushinteger(L, ++i); // add 1 and push it to the top of the stack
     lua_copy(L, -1, lua_upvalueindex(1)); // Update the upvalue count with the new value at the top of the stack
     return 1; // return the count at the top of the stack
}

// factory function, create closure
int new_counter(Lua_State *L) {
     lua_pushinteger(L, 0); // push onto the stack
    
     // Create a C closure, the function prototype is counter, and also
     // includes 1 upvalue, which is 0 pushed in the previous line.
     lua_pushcclosure(L, &counter, 1);

     // The created C closure is pressed on the top of the stack, and the
     // following return 1 means return to the C closure on the top of the stack
     return 1;
}

Let's look at the second function new_counter() first, which is also a factory function for creating closures. First call lua_pushinteger() to push the upvalue count to the top of the stack; then call lua_pushcclosure() to create a closure. To review, a closure consists of a function prototype and some upvalues, which are specified by the last two parameters of the lua_pushcclosure() function. The first parameter specifies the function prototype counter, and the second parameter 1 means that the 1 value at the top of the stack is an upvalue, that is, the 0 just pushed. The following figure is a schematic diagram of the stack before and after calling this function to create a C closure:

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

The far left of the above figure is to push the count i=0 to the top of the stack. In the middle is the created C closure, including the function prototype and upvalue. On the far right is the stack layout after the closure is created, and the closure is pushed onto the stack.

Look at the first function counter() in the above code, which is the function prototype of the closure created. This function is relatively simple, the most critical of which is the lua_upvalueindex() API, which generates an index representing the upvalue, which can be used to read and write the upvalue encapsulated in the closure.

Through the call flow of the code in the above example to the relevant API, we can basically guess the specific implementation of the C closure. Our Rust closures can also refer to this approach. However, Rust natively supports closures! So we can use this feature to implement Rust closures in Lua more simply.

Rust Closure Definition

To implement the "Rust closure" type in Lua with the closure of the Rust language, it is to create a new Value type including the closure of the Rust language.

"Rust Programming Language" has introduced Rust's closures in detail, so I won't say more here. We just need to know that Rust closures are a trait. Specifically, the Rust closure type in Lua is FnMut (&mut ExeState) -> i32. Then you can try to define the Rust closure type of Value in Lua as follows:

pub enum Value {
     RustFunction(fn (&mut ExeState) -> i32), // normal function
     RustClosure(FnMut (&mut ExeState) -> i32), // Closure

However, this definition is illegal, and the compiler will report the following error:

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

This involves the static dispatch and dynamic dispatch of traits in Rust. "Rust Programming Language" also has a detailed introduction for this, so I won’t say more here.

Then, we add dyn according to the compiler's prompt:

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

The compiler still reports an error, but with a different one:

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

That is to say, the trait object is a DST. This was introduced in Introduction to String Definition before, but what I encountered at that time was slice, and now it is trait, which are also the two most important DSTs in Rust. "Rust Programming Language" also has a detailed introduction for this. The solution is to encapsulate a layer of pointers outside. Since Value supports Clone, Box cannot be used, only Rc can be used. And because it is FnMut instead of Fn, it will change the captured environment when it is called, so another layer of RefCell is needed to provide internal variability. So the following definition is obtained:

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

Finally compiled this time! However, think about why Rc<str> was not used when introducing various definitions of strings? Because for the DST type, the actual length needs to be stored in the external pointer or reference, then the pointer will become a "fat pointer", which needs to occupy 2 words. This will further cause the size of the entire Value to become larger. In order to avoid this situation, we can only add another layer of Box, let the Box contain a specific length and become a fat pointer, so that Rc can restore 1 word. It is defined as follows:

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

After defining the type of Rust closure, I also encountered the same problem as Lua closure: should I keep the type of Rust function? It doesn't make much difference whether to keep it or not. We chose to keep it here.

Virtual Machine Execution

The virtual machine implementation of Rust closures is very simple. Because closures and functions are called in the same way in the Rust language, the invocation of Rust closures is the same as the invocation of previous Rust functions:

     fn do_call_function(&mut self, narg_plus: u8) -> usize {
         match self.stack[self.base - 1].clone() {
             Value::RustFunction(f) => { // Rust normal function
                 // omit the preparation of parameters
                 f(self) as usize
             }
             Value::RustClosure(c) => { // Rust closure
                 // Omit the same parameter preparation process
                 c.borrow_mut()(self) as usize
             }

Test

This completes the Rust closure type. After borrowing the closure of the Rust language itself, this implementation is very simple. There is no need to use the Lua stack to cooperate like the official Lua implementation, and there is no need to introduce some special APIs.

The following code shows how the counter example at the beginning of this section can be implemented using Rust closures:

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
}

Compared with the C closure at the beginning of this section, this version is clearer except for the last statement to create a closure, which is very verbose. The last statement will also be optimized when the interpreter API is sorted out later.

Limitations of Rust Closures

As you can see from the sample code above, the captured environment (or upvalue) i needs to be moved into the closure. This also leads to the fact that upvalues cannot be shared among multiple closures. Lua's official C closure does not support sharing, so there is no problem.

Another point that needs to be explained is that Lua's official C closure uses Lua's stack to store upvalue, which leads to the type of upvalue being Lua's Value type. And we use the closure of Rust language, then upvalue can be "more" types, not limited to the Value type. However, the two should be functionally equivalent:

The "more" types supported by Rust closures can be implemented in Lua with LightUserData, that is, pointers; although this is very unsafe for Rust.
The internal types supported in Lua, such as Table, can also be obtained through the API get() in our interpreter (and In Lua's official implementation, the table type is internal and not external).

Build a Lua Interpreter in Rust