From cd57539ce802b3bc2ef896c0b8221d616e8e5d96 Mon Sep 17 00:00:00 2001 From: Maarten Zuidhoorn Date: Wed, 13 Dec 2023 10:35:20 +0100 Subject: [PATCH] BREAKING: Synchronously initialise WASM modules (#2024) When the `experimental.wasm` flag in the Snaps CLI is enabled, the WASM module will now be synchronously initialised, rather than being inlined as a `Uint8Array`. For example, if the WASM module exports a `fibonacci` function, the following is now possible: ```ts import { fibonacci } from './path/to/file.wasm'; fibonacci(...); ``` This is a breaking change, because previously the WASM module would be inlined as `Uint8Array` instead. --- .../examples/packages/wasm/.depcheckrc.json | 1 + .../packages/wasm/assembly/program.ts | 18 ++- .../examples/packages/wasm/snap.manifest.json | 2 +- .../examples/packages/wasm/src/bindings.ts | 12 ++ packages/examples/packages/wasm/src/index.ts | 46 +------ .../examples/packages/wasm/src/types.d.ts | 11 +- packages/snaps-cli/README.md | 2 +- .../webpack/loaders/__fixtures__/program.wasm | Bin 0 -> 139 bytes .../src/webpack/loaders/wasm.test.ts | 95 ++++++++++++- .../snaps-cli/src/webpack/loaders/wasm.ts | 126 +++++++++++++++--- 10 files changed, 246 insertions(+), 67 deletions(-) create mode 100644 packages/examples/packages/wasm/src/bindings.ts create mode 100644 packages/snaps-cli/src/webpack/loaders/__fixtures__/program.wasm diff --git a/packages/examples/packages/wasm/.depcheckrc.json b/packages/examples/packages/wasm/.depcheckrc.json index 15d64e734b..3962c7ec17 100644 --- a/packages/examples/packages/wasm/.depcheckrc.json +++ b/packages/examples/packages/wasm/.depcheckrc.json @@ -5,6 +5,7 @@ "@lavamoat/preinstall-always-fail", "@metamask/auto-changelog", "@metamask/eslint-*", + "@metamask/wasm-example-snap", "@types/*", "@typescript-eslint/*", "eslint-config-*", diff --git a/packages/examples/packages/wasm/assembly/program.ts b/packages/examples/packages/wasm/assembly/program.ts index bf121396bf..2f1eb308e9 100644 --- a/packages/examples/packages/wasm/assembly/program.ts +++ b/packages/examples/packages/wasm/assembly/program.ts @@ -1,3 +1,19 @@ +/** + * Add two numbers together. This is an external function, that will be + * imported from the WebAssembly module, to demonstrate how to import functions + * from the WebAssembly module. + * + * We must use `@external` here to specify the path to the file that contains + * the function. + * + * @param a - The first number. + * @param b - The second number. + * @returns The sum of `a` and `b`. + */ +// eslint-disable-next-line prettier/prettier +@external("../src/bindings.ts", "add") +declare function add(a: i32, b: i32): i32; + /** * Get a fibonacci number after `n` iterations (Fₙ), starting from 1. This is a * TypeScript function that will be compiled to WebAssembly, using @@ -17,7 +33,7 @@ export function fibonacci(iterations: i32): i32 { if (iterations > 0) { // eslint-disable-next-line no-param-reassign, no-plusplus while (--iterations) { - const total = first + second; + const total = add(first, second); first = second; second = total; } diff --git a/packages/examples/packages/wasm/snap.manifest.json b/packages/examples/packages/wasm/snap.manifest.json index 2c89c6550a..298c8db210 100644 --- a/packages/examples/packages/wasm/snap.manifest.json +++ b/packages/examples/packages/wasm/snap.manifest.json @@ -7,7 +7,7 @@ "url": "https://github.com/MetaMask/snaps.git" }, "source": { - "shasum": "Ur4GT3eiORpx5oBr1uXjEwpuuq6nKRM/l6i3grQngWY=", + "shasum": "169i75gNQVsT6z0PTUI5t+KxBhr0fB6APA8J9ZMMm0o=", "location": { "npm": { "filePath": "dist/bundle.js", diff --git a/packages/examples/packages/wasm/src/bindings.ts b/packages/examples/packages/wasm/src/bindings.ts new file mode 100644 index 0000000000..c8b1aaca35 --- /dev/null +++ b/packages/examples/packages/wasm/src/bindings.ts @@ -0,0 +1,12 @@ +/** + * Add two numbers together. This function is not used in the Snap, but is + * imported from the WebAssembly module, to demonstrate how to import functions + * from the WebAssembly module. + * + * @param a - The first number. + * @param b - The second number. + * @returns The sum of `a` and `b`. + */ +export function add(a: number, b: number) { + return a + b; +} diff --git a/packages/examples/packages/wasm/src/index.ts b/packages/examples/packages/wasm/src/index.ts index 7f5d170b3c..6ad9e5e6f8 100644 --- a/packages/examples/packages/wasm/src/index.ts +++ b/packages/examples/packages/wasm/src/index.ts @@ -1,13 +1,14 @@ import { rpcErrors } from '@metamask/rpc-errors'; import type { OnRpcRequestHandler } from '@metamask/snaps-sdk'; +// This is only imported for its type. It is not used at runtime. // eslint-disable-next-line import/order -import { instantiate } from '../build/program'; +import type { instantiate } from '../build/program'; -// This is the WASM module, generated by AssemblyScript, inlined as a -// `Uint8Array` by the MetaMask Snaps CLI loader. +// This is the WASM module, generated by AssemblyScript, inlined as an object +// containing the functions exported by the WASM module. // eslint-disable-next-line import/extensions -import program from '../build/program.wasm'; +import * as program from '../build/program.wasm'; /** * The type of the WASM module. This is generated by AssemblyScript. @@ -19,34 +20,6 @@ type Program = Awaited>; */ type Method = Exclude; -let wasm: Program; - -/** - * Instantiate the WASM module and store the exports in `wasm`. This function - * is called lazily, when the first JSON-RPC request is received. - * - * We use the WASM loader built-in to the MetaMask Snaps CLI to load the WASM - * module. This loader inlines the WASM module as a `Uint8Array` in the - * generated JavaScript bundle, which we can then pass to the WebAssembly - * `instantiate` function. - * - * For this example, we're using AssemblyScript to generate the WASM module, but - * you can use any language that can compile to WASM. - * - * @returns A promise that resolves when the WASM module is instantiated. - * @throws If the WASM module fails to instantiate. - */ -const initializeWasm = async () => { - try { - const module = await WebAssembly.compile(program); - wasm = await instantiate(module, {}); - } catch (error) { - // eslint-disable-next-line no-console - console.error('Failed to instantiate WebAssembly module.', error); - throw error; - } -}; - /** * Handle incoming JSON-RPC requests from the dapp, sent through the * `wallet_invokeSnap` method. This handler handles a single method: @@ -64,11 +37,6 @@ const initializeWasm = async () => { * @see https://developer.mozilla.org/docs/WebAssembly */ export const onRpcRequest: OnRpcRequestHandler = async ({ request }) => { - // Instantiate the WASM module if it hasn't been instantiated yet. - if (!wasm) { - await initializeWasm(); - } - // For this example, we don't validate the request. We assume that the // request is valid, and that the snap is only called with valid requests. In // a real snap, you should validate the request before calling the WASM @@ -76,8 +44,8 @@ export const onRpcRequest: OnRpcRequestHandler = async ({ request }) => { const method = request.method as Method; const params = request.params as Parameters; - if (wasm[method]) { - return wasm[method](...params); + if (program[method]) { + return program[method](...params); } throw rpcErrors.methodNotFound({ data: { method } }); diff --git a/packages/examples/packages/wasm/src/types.d.ts b/packages/examples/packages/wasm/src/types.d.ts index 4769762776..059ad7b28a 100644 --- a/packages/examples/packages/wasm/src/types.d.ts +++ b/packages/examples/packages/wasm/src/types.d.ts @@ -9,6 +9,13 @@ */ // eslint-disable-next-line import/unambiguous declare module '*.wasm' { - const module: Uint8Array; - export default module; + import type { instantiate } from '@metamask/wasm-example-snap/build/program'; + + /** + * The type of the WASM module. This is generated by AssemblyScript. + */ + type Program = Awaited>; + + const module: Program; + export = module; } diff --git a/packages/snaps-cli/README.md b/packages/snaps-cli/README.md index 6ae8b23a3a..c716737cac 100644 --- a/packages/snaps-cli/README.md +++ b/packages/snaps-cli/README.md @@ -422,7 +422,7 @@ can be imported in the snap, for example: ```typescript import program from './program.wasm'; -const module = await WebAssembly.instantiate(program, {}); +// Program is initialised synchronously. // ... ``` diff --git a/packages/snaps-cli/src/webpack/loaders/__fixtures__/program.wasm b/packages/snaps-cli/src/webpack/loaders/__fixtures__/program.wasm new file mode 100644 index 0000000000000000000000000000000000000000..c28966f5ed13457949ae7383a819611945c65505 GIT binary patch literal 139 zcmWN|I}XAy6hzUP_W}qU2}D6bWwU^!LDWghD)I$cD2^1F~WdHKr9NND3;C33PdHD$r58kP%rlUoT cXI5#of}b0*%8KO)PA#2Sb { - it('inlines the WASM module as a `Uint8Array`', async () => { - const source = 'WASM module'; - const result = loader(source); +import loader, { getExports, getImports, getModuleImports, raw } from './wasm'; + +describe('getImports', () => { + it('returns the imports code for the WASM module', () => { + const result = getImports({ + '../src/bindings.ts': ['add', 'subtract'], + './foo.ts': ['bar'], + }); + + expect(result).toMatchInlineSnapshot(` + "import { add, subtract } from "../src/bindings.ts"; + import { bar } from "./foo.ts";" + `); + }); +}); + +describe('getModuleImports', () => { + it('returns the imports code for the WASM module', () => { + const result = getModuleImports({ + '../src/bindings.ts': ['add', 'subtract'], + './foo.ts': ['bar'], + }); + + expect(result).toMatchInlineSnapshot(` + ""../src/bindings.ts": { add, subtract }, + "./foo.ts": { bar }," + `); + }); +}); + +describe('getExports', () => { + it('returns the exports code for the WASM module', () => { + const result = getExports([ + { + kind: 'function', + name: 'fibonacci', + }, + { + kind: 'memory', + name: 'memory', + }, + { + kind: 'function', + name: 'default', + }, + ]); - expect(result).toMatchInlineSnapshot( - `"export default new Uint8Array([87,65,83,77,32,109,111,100,117,108,101]);"`, + expect(result).toMatchInlineSnapshot(` + "export const fibonacci = exports["fibonacci"]; + export const memory = exports["memory"]; + export default exports["default"];" + `); + }); +}); + +describe('loader', () => { + it('synchronously initialises the WASM module', async () => { + const source = await readFile( + join(__dirname, '__fixtures__', 'program.wasm'), ); + + // @ts-expect-error - We don't need to mock the entire `this` object. + const result = await loader.bind({ + addDependency: jest.fn(), + // @ts-expect-error - The type of this function seems to be incorrect. + })(source); + + expect(result).toMatchInlineSnapshot(` + " + import { add } from "../src/bindings.ts"; + + const bytes = new Uint8Array([0,97,115,109,1,0,0,0,1,12,2,96,2,127,127,1,127,96,1,127,1,127,2,26,1,18,46,46,47,115,114,99,47,98,105,110,100,105,110,103,115,46,116,115,3,97,100,100,0,0,3,2,1,1,5,3,1,0,0,7,22,2,9,102,105,98,111,110,97,99,99,105,0,1,6,109,101,109,111,114,121,2,0,10,54,1,52,1,3,127,65,1,33,1,32,0,65,0,74,4,64,3,64,32,0,65,1,107,34,0,4,64,32,2,32,1,16,0,33,3,32,1,33,2,32,3,33,1,12,1,11,11,32,1,15,11,65,0,11]); + const module = new WebAssembly.Module(bytes); + const instance = new WebAssembly.Instance(module, { + "../src/bindings.ts": { add }, + }); + + const exports = instance.exports; + export const fibonacci = exports["fibonacci"]; + export const memory = exports["memory"]; + " + `); + }); + + describe('raw', () => { + it('is `true`', () => { + expect(raw).toBe(true); + }); }); }); diff --git a/packages/snaps-cli/src/webpack/loaders/wasm.ts b/packages/snaps-cli/src/webpack/loaders/wasm.ts index 24a2175e97..75af43650d 100644 --- a/packages/snaps-cli/src/webpack/loaders/wasm.ts +++ b/packages/snaps-cli/src/webpack/loaders/wasm.ts @@ -1,27 +1,121 @@ -import { assert, stringToBytes } from '@metamask/utils'; +/* eslint-disable no-restricted-globals */ + +import { assert } from '@metamask/utils'; +import type { LoaderDefinitionFunction } from 'webpack'; /** - * A Webpack loader that inlines the WASM module as a `Uint8Array`. This makes - * it possible to import the WASM module directly, and use it with the - * `WebAssembly.instantiate` function. + * Get the imports code for the WASM module. This code imports each of the + * imports from the WASM module. * - * This is useful, because snaps are not allowed to import assets from outside - * of their package. This loader allows you to inline the WASM module as a - * `Uint8Array`, which can then be passed to `WebAssembly.instantiate`. + * @param importMap - The import map for the WASM module. + * @returns The imports code for the WASM module. + */ +export function getImports(importMap: Record) { + return Object.entries(importMap) + .map( + ([moduleName, exportNames]) => + `import { ${exportNames.join(', ')} } from ${JSON.stringify( + moduleName, + )};`, + ) + .join('\n'); +} + +/** + * Get the imports code to use in `WebAssembly.Instance`. This code adds each of + * the imports to the `imports` object. * - * @param source - The WASM module as a string. - * @returns A string that exports the WASM module as a `Uint8Array`. + * @param importMap - The import map for the WASM module. + * @returns The imports code for the WASM module. + */ +export function getModuleImports(importMap: Record) { + return Object.entries(importMap) + .map( + ([moduleName, exportNames]) => + `${JSON.stringify(moduleName)}: { ${exportNames.join(', ')} },`, + ) + .join('\n'); +} + +/** + * Get the exports code for the WASM module. This code exports each of the + * exports from the WASM module as a variable. This function assumes that the + * exports are available in a variable named `exports`. + * + * @param descriptors - The export descriptors from the WASM module. + * @returns The exports code for the WASM module. + */ +export function getExports(descriptors: WebAssembly.ModuleExportDescriptor[]) { + return descriptors + .map((descriptor) => { + if (descriptor.name === 'default') { + return `export default exports[${JSON.stringify(descriptor.name)}];`; + } + + return `export const ${descriptor.name} = exports[${JSON.stringify( + descriptor.name, + )}];`; + }) + .join('\n'); +} + +/** + * A Webpack loader that synchronously loads the WASM module. This makes it + * possible to import the WASM module directly. + * + * @param source - The WASM module as `Uint8Array`. + * @returns The WASM module as a JavaScript string. * @example * ```ts - * import wasm from './program.wasm'; + * import * as wasm from './program.wasm'; * - * const { instance } = await WebAssembly.instantiate(wasm, {}); * // Do something with the WASM module... * ``` */ -export default function loader(source: unknown) { - assert(typeof source === 'string', 'Expected source to be a string.'); +// Note: This function needs to be defined like this, so that Webpack can bind +// `this` to the loader context, and TypeScript can infer the type of `this`. +const loader: LoaderDefinitionFunction = async function loader( + source: unknown, +) { + assert(source instanceof Uint8Array, 'Expected source to be a Uint8Array.'); - const bytes = stringToBytes(source); - return `export default new Uint8Array(${JSON.stringify(Array.from(bytes))});`; -} + const bytes = new Uint8Array(source); + const wasmModule = await WebAssembly.compile(bytes); + + // eslint-disable-next-line @typescript-eslint/no-shadow + const exports = WebAssembly.Module.exports(wasmModule); + const imports = WebAssembly.Module.imports(wasmModule).reduce< + Record + >((target, descriptor) => { + target[descriptor.module] ??= []; + target[descriptor.module].push(descriptor.name); + + return target; + }, {}); + + // Add the WASM import as a dependency so that Webpack will watch it for + // changes. + for (const name of Object.keys(imports)) { + this.addDependency(name); + } + + return ` + ${getImports(imports)} + + const bytes = new Uint8Array(${JSON.stringify(Array.from(source))}); + const module = new WebAssembly.Module(bytes); + const instance = new WebAssembly.Instance(module, { + ${getModuleImports(imports)} + }); + + const exports = instance.exports; + ${getExports(exports)} + `; +}; + +export default loader; + +// By setting `raw` to `true`, we are telling Webpack to provide the source as a +// `Uint8Array` instead of converting it to a string. This allows us to avoid +// having to convert the source back to a `Uint8Array` in the loader. +export const raw = true;