Debugging Memory Issues in Rust, WASM, and JavaScript

Lately, I've been experimenting with Rust and WASM and have built a simple implementation of Conway's Game of Life. It's given me an opportunity to better understand some of the pitfalls when working with Rust, WASM, and JavaScript.

I'd like to look at an issue I had with memory while I was working with the project. Let's start with some of the implementation details first.

In Conway's Game of Life you have a grid of "alive" or "dead" creatures (cells). In Rust, you can represented that with this data structure.

#[wasm_bindgen]
pub struct Universe {
    pub width: u32,
    pub height: u32,
    cells: FixedBitSet,
}

The FixedBitSet allows us to represent "alive" and "dead" using single bits.

There are two main functions needed to implement the game of life. The first tells you how many living neighbors are around a particular cell. The second looks at the state of the universe and modifies the cells for the next step in the universe's life cycle (the tick function).

And then there is a function that is intended to exposes the cell data to JavaScript land. On the Universe implementation, you have this function.

impl Universe {
    // Returns a pointer to the memory in self.cells
    pub fn cells_ptr(&self) -> *const u32 {
        self.cells.as_slice().as_ptr()
    }
}

When the Rust code is compiled to WASM using wasm-pack, wasm-pack create a JavaScript module that calls into the WASM. When everything is said and done, I can import both the Universe constructor and also access the WASM memory with my cell data in it.

import { Universe } from "wasm-game-of-life";
import { memory } from "wasm-game-of-life/wasm_game_of_life_bg";

(function() {
    let universe = Universe.new();
    let cellsPtr = universe.cells_ptr();
}());

If we log the cellPtr, we can see that it is a number. That number tells us where in memory our cell data is at. And we can access the data like this...

const cells = new Uint8Array(memory.buffer, cellsPtr, width * height / 8);

Uint8Array is an array of 8-bit unsigned integers. Meaning, each entry in the array contains 8 bits of information. The first argument is an ArrayBuffer, and the memory.buffer variable is an special ArrayBuffer called WebAssembly.Memory. The second argument is the byte offset where we want to start reading from memory. The final argument indicates how many bytes we want to read. Remember how we used FixedBitSet in our Rust code? That data structure truly represents our data as bits. Which means that we only need to read width times height bits of memory from the buffer. If we want that length represented in bytes, then we divide by 8.

So in JavaScript land we have the cell data as an array of bytes (Uint8Array). If we want to get the alive/dead status of a particular cell, then we first need to find the byte that contains our cell. We can do that by dividing the cell location by 8 and rounding down.

// Where n is the position in our universe of cells
const byteIndex = Math.floor(n / 8);

Second, we need to construct a bitmask that will tell us the value of the particular bit within the byte. If we divide the position in the universe by 8 and take the remainder, then we can construct a bitmask shifted by the remainder.

// e.g. 1 << 1 = 10,  1 << 2 = 100
//      1 << 3 = 1000, etc...
const mask = 1 << (n % 8);

Finally, we can do a bitwise AND operation on the mask and the byte, and that will return a byte where the bits are 1 only if it's 1 in both bytes. That means that if our target bit is "on" then our bitwise AND operation should return a value that equals our mask.

const bitIsOn = (cells[byteIndex] & mask) === mask;

So now we can read and display our cells data in the browser, and every quarter of a second call the tick function and re-render the cells.

run(function() {
  const boardElement = document.getElementById('board');
  const universe = Universe.new();
  let cellsPtr = universe.cells_ptr();
  console.log('pointer', cellsPtr);
  boardElement.textContent = renderCells(universe, cellsPtr);
  
  window.setInterval(function() {
    universe.tick();
    console.log('pointer', cellsPtr);
    boardElement.textContent = renderCells(universe, cellsPtr);  
  }, 1000)
})

You'll notice that I'm console logging the pointer and that I'm only fetching the pointer at the start. Which means that I'm assuming my data will always be at the same location in memory.

But consider this implementation of the tick function in Rust - in particular the part at the end...

pub fn tick(&mut self) {
    let mut next = self.cells.clone();

    for row in 0..self.height {
        for col in 0..self.width {
        let idx = self.get_index(row, col);
        let cell = self.cells[idx];
        let live_neighbors = self.live_neighbor_count(row, col);

        let next_cell = match (cell, live_neighbors) {
            // Rule 1: Any live cell with fewer than two live neighbors
            // dies, as if caused by underpopulation.
            (true, x) if x < 2 => false,
            // Rule 2: Any live cell with two or three live neighbors
            // lives on to the next generation.
            (true, 2) | (true, 3) => true,
            // Rule 3: Any live cell with more than three live
            // neighbors dies, as if by overpopulation.
            (true, x) if x > 3 => false,
            // Rule 4: Any dead cell with exactly three live neighbors
            // becomes a live cell, as if by reproduction.
            (false, 3) => true,
            // All other cells remain in the same state.
            (otherwise, _) => otherwise,
        };
        
        next.set(idx, next_cell);
        }
    }

    // Mutation 1 - Points the data to a new location in memory
    // self.cells = next;

    // Mutation 2 - Works - mutates the existing data in memory and
    // keeps the pointer at the same location
    for idx in 0..self.cells.len() {
        self.cells.set(idx, next[idx]);
    }        
}

If I was to mutate my data in-place (Mutation 2), then my previous JavaScript implementation works because my pointer is still pointed to the same memory location. But if I do Mutation 1, I will update where the cell data is pointing to. In that case, I would have to update my JavaScript so it fetches a fresh pointer on every tick.

run(function() {
  const boardElement = document.getElementById('board');
  const universe = Universe.new();
  let cellsPtr = universe.cells_ptr();
  console.log('pointer', cellsPtr);
  boardElement.textContent = renderCells(universe, cellsPtr);
  
  window.setInterval(function() {
    universe.tick();
    // cellsPtr is stale because the tick function
    // has mutated the location of our data.
    //
    // Make sure to get the most recent
    // memory location.
    cellsPtr = universe.cells_ptr();

    console.log('pointer', cellsPtr);
    boardElement.textContent = renderCells(universe, cellsPtr);  
  }, 1000)
})

All that to say, I think this situation highlights some important considerations when you're working with Rust, WASM, and JavaScript. Understanding memory is important, and if you're going to build apps that use Rust, WASM, and JavaScript there are definitely some memory pitfalls to watch out for.