Friday, April 7, 2023

Some considerations for using closures in Rust/WASM

Here are a couple of things I learned while trying to pass a Rust closure to a JavaScript function. Some of these notes are a result of my lack of experience with Rust and Rust/WASM.

Passing a closure that is going to outlive the current function call

Passing a Rust function that exists in the stack to JavaScript is easy for example, here is a call to Array.map:

#[wasm_bindgen]
pub fn square_elements(a : &js_sys::Array) -> js_sys::Array {
    a.map(&mut |value:JsValue, idx: u32, arr: js_sys::Array| {
        let value = value.as_f64().unwrap();
        JsValue::from_f64(value * value)
    })
}

In this case the function used as argument to map exists only for the time of the invocation of the square_elements function. You need to do something extra when passing a closure is going to outlive the current call. This section of the Rust WASM documentation: https://rustwasm.github.io/wasm-bindgen/reference/passing-rust-closures-to-js.html#heap-allocated-closures has details on how to call a JavaScript function that receives a closure that “survives” the current method or function call. A typical example is calling requestAnimationFrame .

It has a very important note that I overlooked at first:

Once a Closure is dropped, it will deallocate its internal memory and invalidate the corresponding JavaScript function so that any further attempts to invoke it raise an exception…https://rustwasm.github.io/wasm-bindgen/reference/passing-rust-closures-to-js.html#heap-allocated-closures

One thing that I want to do in Rust with WASM is to write a “requestAnimationFrame loop” which allows me to write code that performs a repetitive task without blocking the UI thread of the browser. Here is an example of how this looks in JavaScript

// JavaScript
let i = 0;
let action = () => {
   if (i < 3) {
      console.log(`Calling with ${i}`);
      // do something interesting...
      requestAnimationFrame(action);
   }
   i++;
};
requestAnimationFrame(action);

I found a nice example on how to do this loop in Rust here: https://rustwasm.github.io/docs/wasm-bindgen/examples/request-animation-frame.html . The example has a lot of documentation that explains its functionality in the comments. The code looks a little bit intimidating for a Rust newbie (like me). Here is a small reduced code with the most important parts of the example:

let f = Rc::new(RefCell::new(None));
let g = f.clone();

let mut i = 0;
*g.borrow_mut() = Some(Closure::wrap(Box::new(move || {
   if i > 300 {
       ...
       let _ = f.borrow_mut().take(); 
       return;
   }

   i += 1;
   ...
   request_animation_frame(f.borrow().as_ref().unwrap());
}))
...
request_animation_frame(g.borrow().as_ref().unwrap());

As described in the source of the example, this code uses Rc and RefCell to keep the Closure instance alive while the sequence of requestAnimationFrame calls do its work. In this case when the i counter reaches 300 the closure will return and finish the loop.

Each part in this code is very important. I did some mistakes that I’m going to detail in the next sections.

The closure is dropped before ‘requestAnimationFrame’ does its job

The goal of the two Rc/RefCell references to the same closure is to keep the Closure alive before finishing the call to the current function. This is an example of the error that is raised when you fail to do that:

// this example is incomplete
pub fn greet() {
    let f = Rc::new(RefCell::new(None));
    let g = f.clone();
    
    let mut i  = 0;
    *g.borrow_mut() = Some(Closure::new(move || {
         log("At the end of the closure");
    }));
    log(&format!("Before quitting 'greet' {}", Rc::strong_count(&g)));
    request_animation_frame(g.borrow().as_ref().unwrap());
}

Notice that here we don’t pass f into the closure. Hence at the end of the function both g and f are going to be dropped along with the Closure instance. Running this code shows the following errors in the browser console:

...
Before quitting 'greet' 2 wasm_loop_bg.js:259:13
Uncaught Error: closure invoked recursively or after being dropped
...

To fix this issue in the incomplete example I just need to move the f instance to the closure so it will be captured:

// this example is incomplete
pub fn greet() {
    let f = Rc::new(RefCell::new(None));
    let g = f.clone();
    
    let mut i  = 0;
    *g.borrow_mut() = Some(Closure::new(move || {
         let _ = f; /// Now 'f' is moved inside!
         log("At the end of the closure");
    }));
    log(&format!("Before quitting 'greet' {}", Rc::strong_count(&g)));
    request_animation_frame(g.borrow().as_ref().unwrap());
}

Resources not being released

After writing the complete loop, I also found that I was not doing the complete cleanup for the closure.

Here is an example that shows the problematic code:

#[derive(Debug)]
struct MyStruct {
    x: i32
}

impl Drop for MyStruct {
    fn drop(&mut self) {
        write_debug(format!("Calling `drop` on MyStruct: {}", self.x).as_str());
    }
}

#[wasm_bindgen]
pub fn greet() {
    let f = Rc::new(RefCell::new(None));
    let g = f.clone();
    let captured = MyStruct { x: 100 };

    let mut i  = 0;
    *g.borrow_mut() = Some(Closure::new(move || {
        log(format!("Counter: {} captured value: {:?}", i, &captured).as_str());
        if i != 2 {
           request_animation_frame(f.borrow().as_ref().unwrap());
        }
        } else {
            log("Finished");                                                                                                        return;
        }
        i += 1;
    }));

    request_animation_frame(g.borrow().as_ref().unwrap());
    log("Before finishing 'greet'");
}

In this example I created a dummy struct called MyStruct . This structure implements the Drop trait to display a message in the console when the structured is being dropped. An instance of this structure is being captured by the closure passed to the requestAnimationFrame call.

When running this code we see the following messages in the console:

Before finishing 'greet' wasm_loop_bg.js:267:13
Counter: 0 captured value: MyStruct { x: 100 } wasm_loop_bg.js:267:13
Counter: 1 captured value: MyStruct { x: 100 } wasm_loop_bg.js:267:13
Counter: 2 captured value: MyStruct { x: 100 } wasm_loop_bg.js:267:13
Finished

Notice that we don’t see the message logged in the drop method of MyStruct. The reason for this is that I forgot to release the value inside the Rc/RefCell wrapper.

Here is the corrected code:

...
#[wasm_bindgen]
pub fn greet() {
    let f = Rc::new(RefCell::new(None));
    let g = f.clone();
    let captured = MyStruct { x: 100 };

    let mut i  = 0;
    *g.borrow_mut() = Some(Closure::new(move || {
        log(format!("Counter: {} captured value: {:?}", i, &captured).as_str());
        if i != 2 {
           request_animation_frame(f.borrow().as_ref().unwrap());
        } else {
            log("Finished");
            let _ = f.take();
            return;
        }
        i += 1;
    }));

    request_animation_frame(g.borrow().as_ref().unwrap());
    log("Before finishing 'greet'");
}

And now here is the output of the code:

Before finishing 'greet' wasm_loop_bg.js:267:13
Counter: 0 captured value: MyStruct { x: 100 } wasm_loop_bg.js:267:13
Counter: 1 captured value: MyStruct { x: 100 } wasm_loop_bg.js:267:13
Counter: 2 captured value: MyStruct { x: 100 } wasm_loop_bg.js:267:13
Finished wasm_loop_bg.js:267:13
Calling `drop` on MyStruct: 100

As with the original example the take method is used to move the Closure out of the Rc/RefCell reference. The value will be dropped at the end of the call.

Not following the requirement of using FnMut for the closure

This is another case of not following the rules and not taking the time to read the error message. I did a small change to the code as follows:

…

fn my_dummy_function_requiring_move(ms: MyStruct) {
   log(format!("-- {:?}", &ms).as_str());
}
...

#[wasm_bindgen]
pub fn greet() {
    let f = Rc::new(RefCell::new(None));
    let g = f.clone();
    let captured = MyStruct { x: 100 };

    let mut i  = 0;
    *g.borrow_mut() = Some(Closure::new(move || {
        my_dummy_function_requiring_move(captured);
        log(format!("Counter: {} captured value: {:?}", i, &captured).as_str());
        if i != 2 {
           request_animation_frame(f.borrow().as_ref().unwrap());
        } else {
            log("Finished");
            let _ = f.take();
            return;
        }
        i += 1;
    }));

    request_animation_frame(g.borrow().as_ref().unwrap());
    log("Before finishing 'greet'");
}

When introducing this code the compiles shows the following error:

error[E0525]: expected a closure that implements the `FnMut` trait, but this closure only implements `FnOnce`
   --> src\lib.rs:74:41
    |
74  |       *g.borrow_mut() = Some(Closure::new(move || {
    |                              ------------ -^^^^^^
    |                              |            |
    |  ____________________________|____________this closure implements `FnOnce`, not `FnMut`
    | |                            |
    | |                            required by a bound introduced by this call
75  | |
76  | |         my_dummy_function_requiring_move(captured);
    | |                                          -------- closure is `FnOnce` because it moves the variable `captured` out of its environment
77  | |         log(format!("Counter: {} captured value: {:?}", i, &captured).as_str());
...   |
85  | |         i += 1;
86  | |     }));
    | |_____- the requirement to implement `FnMut` derives from here
    |
    = note: required for `[closure@src\lib.rs:74:41: 74:48]` to implement `IntoWasmClosure<dyn FnMut()>`
note: required by a bound in `wasm_bindgen::prelude::Closure::<T>::new`
   --> C:\Users\ldfallasu2\.cargo\registry\src\github.com-1ecc6299db9ec823\wasm-bindgen-0.2.84\src\closure.rs:271:12
    |
271 |         F: IntoWasmClosure<T> + 'static,
    |            ^^^^^^^^^^^^^^^^^^ required by this bound in `wasm_bindgen::prelude::Closure::<T>::new`

It is really impressive to see how to compiler shows the location of the error and the related areas. As the error message says, the problem here is that we are moving a value out of the closure. This fact prevents the compiler from assuming that our closure implements the FnMut trait (More information here https://doc.rust-lang.org/stable/book/ch13-01-closures.html#moving-captured-values-out-of-closures-and-the-fn-traits ).

The solution in this case is simple, since the move was not really required I created an alternative version of the function that do not require a “move”:

fn my_dummy_function_requiring_ref(ms: &MyStruct) {
   write_debug(format!("-- {:?}", ms).as_str());
}

The call is changed to my_dummy_function_requiring_ref(&s); which removes the compilation error.

The examples created in this post use the following utility functions and declarations:

fn window() -> web_sys::Window {
    web_sys::window().expect("not global 'window'")
}

fn request_animation_frame(f: &Closure<dyn FnMut()>) {
    window()
        .request_animation_frame(f.as_ref().unchecked_ref())
        .expect("Call to request animation frame ");
}

#[wasm_bindgen]
extern {
    #[wasm_bindgen(js_namespace = console)]
    fn log(s: &str);
}

Conclusions

I think there are two main conclusions for this post:

  1. Read the documentation carefully!
  2. Take time to read the compiler errors. The Rust compiler team put a lot of effort explain the error and help you locate the origin of the problem.
Posted by at
Labels: , ,
Subscribe to: Posts (Atom)