A few months ago I had to perform extensive work optimizing the jackson-js library. To do so, I basically killed useless calls to function, introduced cache, and so on. And it was enough for my company.

However, I had this question in my head, can I make it even faster? By searching on the web I found that Rust has very good support for WebAssembly (WASM) and that one of the advantages of going to wasm will be the execution speed.

So, I started the migration of jackson-js to wasm. As far as I am today, it is not faster. In fact, I did not expect it to be faster because I knew I would make the migration progressively, and to maybe see speed improvement, I have to minimize calls from JS to Rust (wasm), which is currently absolutely not the case.

Still, I'll detail here my approach and the setup.

You'll find my code here

Mixing up Rust and TypeScript

So, everything starts by having a project written in JavaScript and calling some Rust code compiled to WASM. However, I do not have JavaScript, but TypeScript, because the original project was made in TS.

So, I first had to think about how to mix them up. I started from my jackson-js project, and I created a Rust project with wasm-bindgen as a dependency. I also added some utils dependencies that would help me debug the app and connect to the JavaScript environment. It finally goes like this

1[package]
2name = "jackson-wasm"
3version = "0.1.0"
4edition = "2021"
5
6[features]
7default = ["console_error_panic_hook"]
8
9[lib]
10crate-type = ["cdylib", "rlib"]
11
12[dependencies]
13js-sys = "0.3.69"
14wasm-bindgen="0.2"
15
16console_error_panic_hook = { version = "0.1.7", optional = true }
17
18
19[dependencies.web-sys]
20version = "0.3.69"
21features = [ "console" ]

Then, I modified my original package.json of the TypeScript project to include the built Rust project as part of the TypeScript bundle. And I updated my scripts to include an automatic build of the jackson-wasm project before compiling or testing the TypeScript project.

The jackson-wasm lines are the one of interest

1{
2  "scripts": {
3    "build:rs:node": "cd jackson-wasm && wasm-pack build --target nodejs",
4  },
5  "bundledDependencies": [
6    "jackson-wasm"
7  ],
8  "dependencies": {
9    "lodash.clone": "^4.5.0",
10    "lodash.clonedeep": "^4.5.0",
11    "meriyah": "^4.3.7",
12    "reflect-metadata": "^0.1.13",
13    "jackson-wasm": "file:./jackson-wasm/pkg"
14  }
15}

Migrating JS to Rust

Migrating from JS to Rust is easy. To do so, I decided to have a bottom-up approach, migrating first the leaf methods and moving up to the stack.

For instance, considering this call graph bellow, I first migrate the methodD in rust. Then methodB and methodC together. And I end with methodA.

stateDiagram-v2
    direction LR
    [*] --> methodA
    methodA --> methodB
    methodA --> methodC
    methodB --> methodD
    methodC --> methodD
    methodD --> [*]

This approach allows one to migrate its application progressively and also to keep the TypeScript tests. This last point, tests, is crucial because they ensure that I do not change the behavior of the applications.

Calling Rust from TypeScript

Calling Rust code (WASM) from TypeScript is made super simple thanks to the wasm-bindgen project. To expose a Rust function to TypeScript, I only have to add the #[wasm_bindgen] attributes before the method I want to export. The same attribute before a struct will create a TypeScript class.

For example, I create here the JsonParser in Rust and I expose it to my TypeScript environment with another name ("WASMJsonParser").

1#[wasm_bindgen(js_name = "WASMJsonParser")]
2pub struct JsonParser {
3    get_metadata_cache: HashMap<String, HashMap<String, HashMap<Option<String>, JsValue>>>,
4}

I can then use it directly from my TypeScript code with a classic import statement.

1import { WASMJsonParser } from 'jackson-wasm';

Calling JavaScript from Rust

Calling JavaScript from Rust (WASM) is also made easy by the wasm-bindgen project.

In jackson-js, there is a dependency to the project reflect-metadata that provides a reflection API over TypeScript object. It provides, for example, access to class annotations. To use it in Rust, I first had to declare Reflect has a Rust mod. This step allows one to mimic the JavaScript code aspect. Then, in an extern "C", it is possible to declare the available functions.

1#[allow(non_snake_case)]
2mod Reflect {
3    use wasm_bindgen::prelude::*;
4
5    #[wasm_bindgen]
6    extern "C" {
7        #[wasm_bindgen(js_namespace = Reflect)]
8        pub fn getMetadata(
9            metadataKey: &str,
10            target: &JsValue,
11            property_key: &str,
12        ) -> Option<crate::util::JsonDecoratorOptions>;
13
14        #[wasm_bindgen(js_namespace = Reflect, js_name = getMetadata)]
15        pub fn getMetadata_2(
16            metadataKey: &str,
17            target: &JsValue,
18        ) -> Option<crate::util::JsonDecoratorOptions>;
19
20    }
21}

You'll notice that I created twice the method getMetadata. However, to make it work with the same method name (signature) and different numbers of arguments, you must create multiple Rust methods and use the js_name attribute.

Going fast?

Does it go faster after the migration of some methods?

It can, but it is not sure. Each time a method of the Wasm environment is called, memory management have to be performed. This implies some object copy to ensure the proper working of the two environment together. In my case, the all library is not migrated, yet. Thus, there are a lot of copies made and the hybrid app is a bit slower.

However, it is possible to optimize to reduce at most the slowing part. For instance, considering the first approach diagram. When methodB and methodC are migrated, we should modify again methodD and remove the #[wasm_bindgen]. It allows then to perform some optimization such as not cloning object but using object references instead 😄

What's next?

The next step is clear, end the migration, discover more limitations and improvement path, and make this lib even faster (my first improvement was making the app 100 times faster. Maybe we can go further thanks to WASM)