WASM Interpreter

_README.md

Code from my ST6 Tech Talk Jan '20

The talk is inspired from David Beazley PyCon India 2019 Keynote: https://www.youtube.com/watch?v=VUT386_GKI8

Requires:

• Recent version of Node.js and https://github.com/TypeStrong/ts-node
• You will need to apply the patch to the wasm-parser decoder.
• Download the rocket wasm file from: https://github.com/aochagavia/rocket_wasm/blob/master/html/program.wasm

Additional resources:

• Rocket Game in Browser written in Rust: https://aochagavia.github.io/blog/rocket---a-rust-game-running-on-wasm/
• Terminal IO - terminal based game engine: https://github.com/robertrypula/terminal-game-io
• WASM Reference Manual: https://github.com/sunfishcode/wasm-reference-manual/blob/master/WebAssembly.md
• Good intro to WASM: https://rsms.me/wasm-intro
• WASM Interpreter reference implementation: https://github.com/WebAssembly/spec/tree/master/interpreter
• WASM Parser in JavaScript: https://github.com/xtuc/webassemblyjs/tree/master/packages/wasm-parser
• VSCode WebAssembly extension: https://github.com/wasmerio/vscode-wasm
• David Beazley talk's source code: https://gist.github.com/dabeaz/7d8838b54dba5006c58a40fc28da9d5a

game.ts

import * as fs from 'fs';
import { decode } from '@webassemblyjs/wasm-parser';
import { Func, ImportFunc, Machine } from './machine';

const program = decode(fs.readFileSync('./program.wasm'));

const transformInstr = ({ id, args, index, object, instr }) => {
  args = args ? args : (index ? [index] : []);
  id = object ? `${object.replace('u', 'i')}.${id}` : id;

  const toBigInt = ({ low, high }) => {
    if (low < 0) {
      low = 2 ** 32 - low;
    }
    return BigInt(`0x${high.toString(16)}${low.toString(16).padStart(8, '0')}`);
  };

  if (id === 'block' || id === 'loop') {
    args = [, instr.map(i => transformInstr(i))];
  } else {
    args = args.map(a => a.value).filter(v => v != undefined).map(v => isNaN(v) ? toBigInt(v) : v);
  }

  if (id === 'br_table') {
    args = [args.slice(0, args.length - 1), [args[args.length - 1]]];
  }

  return [id, ...args];
}

const definedFunctions = program.body[0].fields
  .filter(f => f.type === 'Func')
  .map(f => new Func(
    f.signature.params.length,
    f.signature.results.length > 0,
    f.body.map(i => transformInstr(i)))
  );

let board = [];

const FPS = 60;
const BOARD_WIDTH = 60;
const BOARD_HEIGHT = 40;

const draw = (x, y, char, z = undefined) => {
  if (z) console.log(z);
  if (0 < x && 0 < y) {
    x = Math.round(x) % BOARD_WIDTH;
    y = Math.round(y) % BOARD_HEIGHT;
    board[x][y] = char;
  }
}

const getPlayer = (radians) => {
  if (0 < radians && radians < Math.PI) {
    return '⬇';
  }
  return '⬆';
}

const implementations = {
  'Math_atan': (x) => Math.atan(x),
  'cos': (x) => Math.cos(x),
  'sin': (x) => Math.sin(x),
  'clear_screen': (x) => { }/* console.log('clear_screen') */,
  'draw_score': (s) => { }/* console.log('draw_score', s) */,

  'draw_bullet': (x, y) => draw(x, y, '❂'),
  'draw_enemy': (x, y) => draw(x, y, '⏣'),
  'draw_particle': (x, y, z) => draw(x, y, z > 2 ? '◦' : z > 1 ? '•' : '∙'),
  'draw_player': (x, y, z) => draw(x, y, getPlayer(z)),

  // 'draw_bullet': (x, y) => console.log('draw_bullet', x, y),
  // 'draw_enemy': (x, y) => console.log('draw_enemy', x, y),
  // 'draw_particle': (x, y, z) => console.log('draw_particle', x, y, z),
  // 'draw_player': (x, y, z) => console.log('draw_player', x, y, z),


};

const importedFunctions = [
  new ImportFunc(1, 1, implementations.Math_atan),
  new ImportFunc(0, 0, implementations.clear_screen),
  new ImportFunc(1, 1, implementations.cos),
  new ImportFunc(2, 0, implementations.draw_bullet),
  new ImportFunc(2, 0, implementations.draw_enemy),
  new ImportFunc(3, 0, implementations.draw_particle),
  new ImportFunc(3, 0, implementations.draw_player),
  new ImportFunc(1, 0, implementations.draw_score),
  new ImportFunc(1, 1, implementations.sin),
];

const functions = [...importedFunctions, ...definedFunctions];


const machine = new Machine(functions, 20 * 65536);

// Initialize memory
const data = program.body[0].fields.filter(f => f.type === 'Data').map(f => ({
  op: transformInstr(f.offset),
  values: f.init.values,
}));

const memoryBytes = new Uint8Array(machine.memory.buffer);

for (const { op, values } of data) {
  machine.execute([op]);
  const offset = machine.pop();
  memoryBytes.set(values, offset);
}

machine.call(definedFunctions[11], [BOARD_WIDTH, BOARD_HEIGHT]);

let shoot = 0;
let left = 0;
let right = 0;
let boost = 0;

const frameHandler = (instance) => {
  for (let x = 0; x < BOARD_WIDTH; x++) {
    board[x] = [];
  }

  machine.call(definedFunctions[13], [1 / (FPS * 8)]);
  machine.call(definedFunctions[12], []);

  let frameData = '';

  for (let y = 0; y < BOARD_HEIGHT; y++) {
    for (let x = 0; x < BOARD_WIDTH; x++) {
      frameData += board[x][y] || ' ';
    }
  }
  instance.drawFrame(frameData, BOARD_WIDTH, BOARD_HEIGHT);
};

const TerminalGameIo = require('terminal-game-io');

const Key = TerminalGameIo.Key;

const toggleKeyFunc = (index) => {
  machine.call(definedFunctions[index], [1]);
  setTimeout(() => machine.call(definedFunctions[index], [0]), 500)
}

const keypressHandler = (instance, keyName) => {
  switch (keyName) {
    case Key.Space:
      toggleKeyFunc(14);
      break;
    case Key.ArrowUp:
      toggleKeyFunc(15);
      break;
    case Key.ArrowLeft:
      toggleKeyFunc(16);
      break;
    case Key.ArrowRight:
      toggleKeyFunc(17);
      break;
    case Key.Escape:
      instance.exit();
      break;
  }

  frameHandler(instance);
};

const terminalGameIo = TerminalGameIo.createTerminalGameIo({
  fps: FPS,
  frameHandler,
  keypressHandler
});

machine.ts

import * as util from 'util';

const consts = {
  'i32.const': Number,
  'i64.const': BigInt,
  'f32.const': Number,
  'f64.const': Number,
}

const overflow = (x, n) => {
  const result = x % BigInt(2 ** n);
  if (result >= BigInt(2 ** (n - 1))) {
    return result - BigInt(2 ** n);
  } else if (result < -BigInt(2 ** (n - 1))) {
    return BigInt(2 ** n) + result;
  }
  return result;
}

const i32Unsign = (x) =>
  new DataView(new Int32Array([x]).buffer).getUint32(0, true);

const binary = {
  'i32.add': (x, y) => x + y,
  'i32.sub': (x, y) => x - y,
  'i32.mul': (x, y) => x * y,
  'i32.div_s': (x, y) => Math.round(x / y),
  'i32.div_u': (x, y) => Math.round(x / y),
  'i32.and': (x, y) => x & y,
  'i32.or': (x, y) => x | y,
  'i32.xor': (x, y) => x ^ y,
  'i32.rem_s': (x, y) => x % y,
  'i32.rem_u': (x, y) => x % y,
  'i32.eq': (x, y) => Number(x === y),
  'i32.ne': (x, y) => Number(x != y),
  'i32.lt_s': (x, y) => Number(x < y),
  'i32.le_s': (x, y) => Number(x <= y),
  'i32.gt_s': (x, y) => Number(x > y),
  'i32.ge_s': (x, y) => Number(x >= y),
  'i32.lt_u': (x, y) => Number(i32Unsign(x) < i32Unsign(y)),
  'i32.le_u': (x, y) => Number(i32Unsign(x) <= i32Unsign(y)),
  'i32.gt_u': (x, y) => Number(i32Unsign(x) > i32Unsign(y)),
  'i32.ge_u': (x, y) => Number(i32Unsign(x) >= i32Unsign(y)),
  'i32.rotr': (x, y) => (x >> y) | ((x & ((2 ** y) - 1)) << (32 - y)),
  'i32.rotl': (x, y) => (x << y) | ((x & ((2 ** y) - 1)) >> (32 - y)),
  'i32.shr_u': (x, y) => x >> y,
  'i32.shl': (x, y) => x << y,

  'i64.add': (x, y) => BigInt(x) + BigInt(y) % BigInt(2 ** 64),
  'i64.sub': (x, y) => BigInt(x) - BigInt(y),
  'i64.mul': (x, y) => overflow(BigInt(x) * BigInt(y), 64),
  'i64.div_s': (x, y) => Math.round(x / y),
  'i64.div_u': (x, y) => Math.round(x / y),
  'i64.and': (x, y) => BigInt(x) & BigInt(y),
  'i64.or': (x, y) => BigInt(x) | BigInt(y),
  'i64.xor': (x, y) => BigInt(x) ^ BigInt(y),
  'i64.rem_s': (x, y) => x % y,
  'i64.rem_u': (x, y) => x % y,
  'i64.eq': (x, y) => x === y,
  'i64.ne': (x, y) => x != y,
  'i64.lt_s': (x, y) => x < y,
  'i64.le_s': (x, y) => x <= y,
  'i64.gt_s': (x, y) => x > y,
  'i64.ge_s': (x, y) => x >= y,
  'i64.lt_u': (x, y) => x < y,
  'i64.le_u': (x, y) => x <= y,
  'i64.gt_u': (x, y) => x > y,
  'i64.ge_u': (x, y) => x >= y,
  'i64.rotr': (x, y) => (x >> y) | ((x & ((2 ** y) - 1)) << (64 - y)),
  'i64.rotl': (x, y) => (x << y) | ((x & ((2 ** y) - 1)) >> (64 - y)),
  'i64.shr_u': (x, y) => new DataView(new BigInt64Array([x]).buffer).getBigUint64(0, true) >> BigInt(y),
  'i64.shl': (x, y) => overflow(BigInt(x) << BigInt(y), 64),

  'f64.add': (x, y) => x + y,
  'f64.sub': (x, y) => x - y,
  'f64.mul': (x, y) => x * y,
  'f64.div': (x, y) => x / y,
  'f64.eq': (x, y) => Number(x === y),
  'f64.ne': (x, y) => Number(x != y),
  'f64.lt': (x, y) => Number(x < y),
  'f64.le': (x, y) => Number(x <= y),
  'f64.gt': (x, y) => Number(x > y),
  'f64.ge': (x, y) => Number(x >= y),
};

const unary = {
  'i32.eqz': (x) => Number(x === 0),
  'i32.clz': x => Math.clz32(x),
  'i32.ctz': x => 31 - Math.clz32(x ^ (x - 1)),
  'i64.extend_u/i32': x => BigInt(x >>> 0),
  'f64.reinterpret/i64': x => new DataView(new BigUint64Array([x]).buffer).getFloat64(0, true),
  'f64.sqrt': x => Math.sqrt(x),
  'f64.convert_u/i32': i32Unsign,
  'i32.wrap/i64': x => Number(overflow(x, 32)),
}

const load = {
  'i32.load': DataView.prototype.getInt32,
  'i64.load': DataView.prototype.getBigInt64,
  'f64.load': DataView.prototype.getFloat64,
  'i32.load8_s': DataView.prototype.getInt8,
  'i32.load8_u': DataView.prototype.getUint8,
  'i32.load16_s': DataView.prototype.getInt16,
  'i32.load16_u': DataView.prototype.getUint16,
  'i64.load8_s': DataView.prototype.getInt8,
  'i64.load8_u': DataView.prototype.getUint8,
  'i64.load16_s': DataView.prototype.getInt16,
  'i64.load16_u': DataView.prototype.getUint16,
  'i64.load32_s': DataView.prototype.getInt32,
  'i64.load32_u': DataView.prototype.getUint32,
};

const store = {
  'i32.store': DataView.prototype.setInt32,
  'i64.store': DataView.prototype.setBigInt64,
  'f64.store': DataView.prototype.setFloat64,
  'i32.store8': DataView.prototype.setInt8,
  'i32.store16': DataView.prototype.setInt16,
  'i64.store8': DataView.prototype.setInt8,
  'i64.store16': DataView.prototype.setInt16,
  'i64.store32': DataView.prototype.setInt32,
};

export class Func {
  constructor(public params, public returns, public code) {
  }
}

export class ImportFunc {
  constructor(public params, public returns, public call: Function) {
  }
}

class Break {
  constructor(public level) {
  }
}

class Return {

}

export class Machine {
  public memory: DataView;
  private readonly items: Array<any>;

  constructor(private functions: Array<Func | ImportFunc> = [], memSize = 65536) {
    this.memory = new DataView(new ArrayBuffer(memSize));
    this.items = []
  }

  push(item) {
    this.items.push(item);
  }

  pop() {
    return this.items.pop();
  }

  call(func: Func | ImportFunc, args) {
    if (func instanceof Func) {
      try {
        this.execute(func.code, args);
      } catch (err) {
        if (!(err instanceof Return)) {
          throw err;
        }
      }
      if (func.returns) {
        return this.pop();
      }
    } else {
      // console.dir({ func, args });
      return func.call.apply(null, args);
    }
  }

  formatArray(arr) {
    return '[' + arr.join(', ') + ']';
  }

  execute(ops, locals = []) {
    for (const [op, ...rest] of ops) {
      // if (this.items.length > 0) {
      //   console.log(op, this.formatArray(rest), this.formatArray(this.items));
      // } else {
      //   console.log(op);
      // }
      if (consts[op]) {
        this.push(consts[op](rest[0]));
      } else if (binary[op]) {
        const right = this.pop();
        const left = this.pop();
        this.push(binary[op](left, right));
      } else if (unary[op]) {
        const value = this.pop();
        this.push(unary[op](value));
      }
      else if (load[op]) {
        const [, offset] = rest;
        const addr = this.pop() + offset;
        this.push(load[op].call(this.memory, addr, true));
      } else if (store[op]) {
        let value = this.pop();
        if (op === 'i64.store') {
          value = BigInt(value);
        }
        const [, offset] = rest;
        const addr = this.pop() + offset;
        store[op].call(this.memory, addr, value, true);
      }

      else if (op === 'memory.size' || op === 'current_memory') {
        this.push(Math.round(this.memory.byteLength / 65536));
      } else if (op === 'memory.grow' || op === 'grow_memory') {
        const pages = this.pop();

        const newBuffer = new ArrayBuffer(this.memory.byteLength + pages * 65536);
        const newMemory = new Uint8Array(newBuffer);
        newMemory.set(new Uint8Array(this.memory.buffer))
        this.memory = new DataView(newMemory.buffer);
      }


      else if (op === 'get_local') {
        this.push(locals[rest[0]]);
      } else if (op === 'set_local') {
        locals[rest[0]] = this.pop();
      } else if (op === 'tee_local') {
        locals[rest[0]] = this.items[this.items.length - 1];
      }

      else if (op === 'drop') {
        this.pop();
      }

      else if (op === 'select') {
        const cond = this.pop();
        const v2 = this.pop();
        const v1 = this.pop();

        this.push(cond ? v1 : v2);
      }


      else if (op === 'call') {
        // console.dir({ funcs: this.functions, args: rest })
        const func = this.functions[rest[0]];
        const args = Array.from(Array(func.params).keys()).map(p => this.pop()).reverse();
        const result = this.call(func, args);
        if (func.returns) {
          this.push(result);
        }
      } else if (op === 'br') {
        throw new Break(rest[0]);
      } else if (op === 'br_if') {
        if (this.pop()) {
          throw new Break(rest[0]);
        }
      } else if (op === 'br_table') {
        const n = this.pop();
        if (n < rest[0].length) {
          throw new Break(rest[0][n]);
        } else {
          throw new Break(rest[1]);
        }
      } else if (op === 'block') {
        try {
          this.execute(rest[1], locals);
        } catch (err) {
          if (err instanceof Break) {
            if (err.level > 0) {
              err.level--;
              throw err;
            }
          } else {
            throw err;
          }
        }
      } else if (op === 'loop') {
        while (true) {
          try {
            this.execute(rest[1], locals);
            break;
          } catch (err) {
            if (err instanceof Break) {
              if (err.level > 0) {
                err.level--;
                throw err;
              }
            } else {
              throw err;
            }
          }
        }
      } else if (op === 'return') {
        throw new Return();
      } else if (op === 'local' || op === 'unreachable' || op === 'end' || op === 'call_indirect') {
        // enjoy
      } else {
        console.error('=== Unknown op ===', op);
      }
    }
  }
}

package.json

{
  "dependencies": {
    "@types/node": "13.1.8",
    "@webassemblyjs/wasm-parser": "1.8.5",
    "terminal-game-io": "3.1.0",
    "ts-node": "8.6.2",
    "typescript": "3.7.5"
  }
}

rocket.ts

import * as fs from 'fs';
import { decode } from '@webassemblyjs/wasm-parser';
import { Func, ImportFunc, Machine } from './machine';

const implementations = {
  'Math_atan': (x) => Math.atan(x),
  'clear_screen': (x) => console.log('clear_screen'),
  'cos': (x) => Math.cos(x),
  'draw_bullet': (x, y) => console.log('draw_bullet', x, y),
  'draw_enemy': (x, y) => console.log('draw_enemy', x, y),
  'draw_particle': (x, y, z) => console.log('draw_particle', x, y, z),
  'draw_player': (x, y, z) => console.log('draw_player', x, y, z),
  'draw_score': (s) => console.log('draw_score', s),
  'sin': (x) => Math.sin(x),
};

const importedFunctions = [
  new ImportFunc(1, 1, implementations.Math_atan),
  new ImportFunc(0, 0, implementations.clear_screen),
  new ImportFunc(1, 1, implementations.cos),
  new ImportFunc(2, 0, implementations.draw_bullet),
  new ImportFunc(2, 0, implementations.draw_enemy),
  new ImportFunc(3, 0, implementations.draw_particle),
  new ImportFunc(3, 0, implementations.draw_player),
  new ImportFunc(1, 0, implementations.draw_score),
  new ImportFunc(1, 1, implementations.sin),
];

const transformInstr = ({ id, args, index, object, instr }) => {
  args = args ? args : (index ? [index] : []);
  id = object ? `${object.replace('u', 'i')}.${id}` : id;

  const toBigInt = ({ low, high }) => {
    if (low < 0) {
      low = 2 ** 32 - low;
    }
    return BigInt(`0x${high.toString(16)}${low.toString(16).padStart(8, '0')}`);
  };

  if (id === 'block' || id === 'loop') {
    args = [, instr.map(i => transformInstr(i))];
  } else {
    args = args.map(a => a.value).filter(v => v != undefined).map(v => isNaN(v) ? toBigInt(v) : v);
  }

  if (id === 'br_table') {
    args = [args.slice(0, args.length - 1), [args[args.length - 1]]];
  }

  return [id, ...args];
}

const program = decode(fs.readFileSync('./program.wasm'));

const definedFunctions = program.body[0].fields
  .filter(f => f.type === 'Func')
  .map(f => new Func(
    f.signature.params.length,
    f.signature.results.length > 0,
    f.body.map(i => transformInstr(i)))
  );

const functions = [...importedFunctions, ...definedFunctions];


const machine = new Machine(functions, 20 * 65536);

// Initialize memory
const data = program.body[0].fields.filter(f => f.type === 'Data').map(f => ({
  op: transformInstr(f.offset),
  values: f.init.values,
}));

const memoryBytes = new Uint8Array(machine.memory.buffer);

for (const { op, values } of data) {
  machine.execute([op]);
  const offset = machine.pop();
  memoryBytes.set(values, offset);
}

// Game on
const width = 80.0;
const height = 60.0;

machine.call(definedFunctions[11], [width, height]);

let last = Date.now() / 1000;
while (true) {
  const now = Date.now() / 1000;
  const dt = now - last;
  last = now;

  machine.call(definedFunctions[13], [dt]);
  machine.call(definedFunctions[12], []);;
}

wasm-parser.patch

diff --git a/node_modules/@webassemblyjs/wasm-parser/lib/decoder.js b/node_modules/@webassemblyjs/wasm-parser/lib/decoder.js
index 79e1ee2..af392b3 100644
--- a/node_modules/@webassemblyjs/wasm-parser/lib/decoder.js
+++ b/node_modules/@webassemblyjs/wasm-parser/lib/decoder.js
@@ -837,10 +837,12 @@ function decode(ab, opts) {
           var align = aligun32.value;
           eatBytes(aligun32.nextIndex);
           dump([align], "align");
+          args.push(t.numberLiteralFromRaw(align));
           var offsetu32 = readU32();
           var _offset2 = offsetu32.value;
           eatBytes(offsetu32.nextIndex);
           dump([_offset2], "offset");
+          args.push(t.numberLiteralFromRaw(_offset2));
         }
       } else if (instructionByte >= 0x41 && instructionByte <= 0x44) {
         /**