Rust: Can a WebAssembly::Memory SharedArrayBuffer be read as a slice without copying?

Context

I wrote a small genetics library in JavaScript for a browser-based game. The library consists of entirely static methods that operate on typed arrays. For example:

// 2 implies there are 2 variants of this gene (in our case, "A" or "O"):
const a_type_blood = Genetics.Gene(2); // Returns: [2]

// same as above ("B" or "O"):
const b_type_blood = Genetics.Gene(2); // Returns: [2]

// pass all genes that belong to a species:
const human = Genetics.Species([a_type_blood, b_type_blood]); // Returns: [2, 2]

// randomize the genes for each individual... remember there are 2 copies of each (chromosomes):
const mom = Genetics.RandomGenes(); // Could return, for instance: [1, 1, 1, 2]
const dad = Genetics.RandomGenes(); // Could return, for instance: [2, 2, 2, 1]

// The first 2 numbers are the alleles for gene 1 (e.g. whether they have A type blood)
// The second 2 numbers are the alleles for gene 2 (e.g. whether they have B type blood)
// In this case, the mom has B-type blood (one chromosome produces the B allele)
// The dad has AB-type blood (both chromosomes produce the A allele, one produces the B allele)

// select random chromosomes from each parent:
const child = Genetics.Breed(mom, dad); // Could return, for instance: [1, 2, 1, 1]

// The child got an "O" on the first (A) gene from the mom
// The child got an "A" on the first (A) gene from the dad
// The child got an "O" on the second (B) gene from the mom
// The child got an "O" on the second (B) gene from the dad
// So the child has A blood type

// 3 inputs:
//  - which individual you want the phenotype of
//  - which genes are used to compute the phenotype (in this case, genes 1 and 2)
//  - a function to compute the phenotype from the gene expressions
const bloodType = Genetics.Phenotype(child, [1, 2], function(a, b) {
    if (a && b) return "AB";
    if (a) return "A";
    if (b) return "B";
    return "O";
});

I use this pattern for three reasons:

My library now holds no state. All state is retained by the clients of my library
Typed arrays are trivially serializable, allowing me to easily save species and individuals to LocalStorage or a database
I can take advantage of the fact that arrays in JavaScript are passed by reference and not by value, so performing mutations on the genes is very efficient

Example mutation:

// Remove the A-type alleles from the father
// First gene, first chromosome, new value = 1
Genetics.Splice(dad, 1, 1, 1);
// First gene, second chromosome, new value = 1
Genetics.Splice(dad, 1, 2, 1);

Goal with WASM

Now I would like to utilize this library outside of the browser, in a Unity3D game. So my goal was to:

Re-write the library in Rust
Compile to both a DLL and to WASM, so that the same code can power both Unity3D and the browser

Because WASM runs in a separate thread, ownership of typed arrays must be passed back and forth between JS and WASM. I was concerned about the potential performance cost of this. WASM handily provides a way to access SharedArrayBuffers through WebAssembly.Memory. I did a quick test and was able to read and write to memory much faster like this:

lib.rs:

#[wasm_bindgen]
pub fn read_from_shared_memory(mem: &WebAssembly::Memory) -> JsValue {
    let arr = Uint8Array::new(&mem.buffer());
    return Uint8Array::at(&arr, 0).into();
}

#[wasm_bindgen]
pub fn write_to_shared_memory(mem: &WebAssembly::Memory, val: u8) -> () {
    let arr = Uint8Array::new(&mem.buffer());
    arr.set_index(0, val.into());
}

#[wasm_bindgen]
pub fn read_from_message_passing(arr: &[u8]) -> JsValue {
    return arr[0].into();
}

#[wasm_bindgen]
pub fn write_to_message_passing(arr: &mut [u8], val: u8) -> () {
    arr[0] = val;
}

test.js:

let d0 = new Date();
let sum = 0;
const mem = new WebAssembly.Memory({
  initial: 1, // 1 page = 64k
  maximum: 1,
  shared: true,
});
for (let i = 0; i < 50000; i++) {
  write_to_shared_memory(mem, i);
  sum += read_from_shared_memory(mem);
}
console.log("Value: ", sum, "Using shared memory: ", new Date() - d0);

d0 = new Date();
sum = 0;
let mem2 = new Uint8Array(65536); // 64k
for (let i = 0; i < 50000; i++) {
  write_to_message_passing(mem2, i);
  sum += read_from_message_passing(mem2);
}
console.log("Value: ", sum, "Using message passing: ", new Date() - d0);

Output:

Value:  6367960 Using shared memory:  18
Value:  6367960 Using message passing:  281

That’s a 15.6x speed up by using shared memory!

Issue

Getting this speed up required the use of &WebAssembly::Memory, which is very specific to my WASM build. When I build a DLL for Unity3D, this type will be meaningless. So I need to abstract my data manipulation from the WASM interface. Something like the following:

lib.rs:

#[wasm_bindgen]
pub fn read_from_shared_memory(mem: &WebAssembly::Memory) -> JsValue {
    return read_impl(???);
}

pub fn read_impl(arr: &[u8]) -> u8 {
    return arr[0];
}

#[wasm_bindgen]
pub fn write_to_shared_memory(mem: &WebAssembly::Memory, val: u8) -> () {
    write_impl(???);
}

pub fn write_impl(mem: &mut [u8], val: u8) -> () {
    mem[0] = val;
}

Then from my Unity3D I simply call write_impl but from JavaScript I call write_from_shared_memory

The problem is, I don’t know what the ??? should be!

js_sys handily provides a to_vec method, but this method copies the data to a new location in memory! This not only eliminates my performance win, but also breaks my “write” implementation:

lib.rs:

#[wasm_bindgen]
pub fn read_from_shared_memory(mem: &WebAssembly::Memory) -> JsValue {
    let arr = Uint8Array::new(&mem.buffer());
    return read_impl(&arr.to_vec()[..]).into();
}

#[wasm_bindgen]
pub fn write_to_shared_memory(mem: &WebAssembly::Memory, val: u8) -> () {
    let arr = Uint8Array::new(&mem.buffer());
    write_impl(&mut (arr.to_vec()[..]), val);
}

Output:

// Because we modified the COPY of our shared memory, the values we read back are always 0!
Value:  0 Using shared memory:  298
Value:  6367960 Using message passing:  272

I know that what I’m looking for is unsafe. If a method exists, I expect it to require unsafe Rust in order to execute. But what I want is:

Given a &WebAssembly::Memory
Extract a &mut [u8] pointing to the same location in memory (no copying!)

How do I do this? Can I do this?

Filed under: Kiến thức lập trình - @ 03:46

Thẻ: rustwasm-bindgen

Thiết kế website giá rẻ

Danh mục

Rust: Can a WebAssembly::Memory SharedArrayBuffer be read as a slice without copying?

Context

Goal with WASM

Issue