Skip to content

Commit

Permalink
Allow using a custom loader
Browse files Browse the repository at this point in the history
This deprecates the `targetEnv` option and replaces it with `loader`
which performs the same thing while also accepting an actual function.

When that happens, that function's body will be inserted in lieu of the
default one, giving complete control over the way WASM is loaded into
the runtime.
  • Loading branch information
tiennou committed Feb 20, 2024
1 parent 2a19079 commit 226dfd6
Show file tree
Hide file tree
Showing 4 changed files with 107 additions and 23 deletions.
21 changes: 17 additions & 4 deletions packages/wasm/src/helper.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
import type { TargetEnv } from '../types';
import type { TargetEnv, WasmLoaderFunction } from '../types';

export const HELPERS_ID = '\0wasmHelpers.js';

export const LOADER_FUNC_NAME = '_loadWasmModule';

const nodeFilePath = `
var fs = require("fs")
var path = require("path")
Expand Down Expand Up @@ -94,8 +96,8 @@ const envModule = (env: TargetEnv) => {
}
};

export const getHelpersModule = (env: TargetEnv) => `
function _loadWasmModule (sync, filepath, src, imports) {
const defaultLoader = (env: TargetEnv) => `
function ${LOADER_FUNC_NAME} (sync, filepath, src, imports) {
function _instantiateOrCompile(source, imports, stream) {
var instantiateFunc = stream ? WebAssembly.instantiateStreaming : WebAssembly.instantiate;
var compileFunc = stream ? WebAssembly.compileStreaming : WebAssembly.compile;
Expand All @@ -116,5 +118,16 @@ function _loadWasmModule (sync, filepath, src, imports) {
return _instantiateOrCompile(buf, imports, false)
}
}
export { _loadWasmModule };
`;

export const getHelpersModule = (loader: WasmLoaderFunction | TargetEnv) => {
let code = '';
if (loader instanceof Function) {
code += loader.toString().replace(loader.name, LOADER_FUNC_NAME);
} else {
code += defaultLoader(loader);
}

code += `export { ${LOADER_FUNC_NAME} };`;
return code;
};
19 changes: 12 additions & 7 deletions packages/wasm/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,15 +7,17 @@ import { createFilter } from '@rollup/pluginutils';

import type { RollupWasmOptions } from '../types';

import { getHelpersModule, HELPERS_ID } from './helper';
import { getHelpersModule, HELPERS_ID, LOADER_FUNC_NAME } from './helper';

export function wasm(options: RollupWasmOptions = {}): Plugin {
// eslint-disable-next-line no-param-reassign
options.loader ??= options.targetEnv;
const {
sync = [],
maxFileSize = 14 * 1024,
publicPath = '',
targetEnv = 'auto',
fileName = '[hash][extname]'
fileName = '[hash][extname]',
loader = 'auto'
} = options;

const syncFiles = sync.map((x) => path.resolve(x));
Expand All @@ -35,7 +37,7 @@ export function wasm(options: RollupWasmOptions = {}): Plugin {

load(id) {
if (id === HELPERS_ID) {
return getHelpersModule(targetEnv);
return getHelpersModule(loader);
}

if (!filter(id)) {
Expand All @@ -49,7 +51,7 @@ export function wasm(options: RollupWasmOptions = {}): Plugin {

return Promise.all([fs.promises.stat(id), fs.promises.readFile(id)]).then(
([stats, buffer]) => {
if (targetEnv === 'auto-inline') {
if (loader === 'auto-inline') {
return buffer.toString('binary');
}

Expand Down Expand Up @@ -88,6 +90,7 @@ export function wasm(options: RollupWasmOptions = {}): Plugin {
if (code && /\.wasm$/.test(id)) {
const isSync = syncFiles.indexOf(id) !== -1;
const publicFilepath = copies[id] ? `'${copies[id].publicFilepath}'` : null;
let out = '';
let src;

if (publicFilepath === null) {
Expand All @@ -100,12 +103,14 @@ export function wasm(options: RollupWasmOptions = {}): Plugin {
src = null;
}

out = `import { ${LOADER_FUNC_NAME} } from ${JSON.stringify(HELPERS_ID)};
export default function (imports) { return ${LOADER_FUNC_NAME}(${+isSync}, ${publicFilepath}, ${src}, imports) }`;

return {
map: {
mappings: ''
},
code: `import { _loadWasmModule } from ${JSON.stringify(HELPERS_ID)};
export default function(imports){return _loadWasmModule(${+isSync}, ${publicFilepath}, ${src}, imports)}`
code: out
};
}
return null;
Expand Down
59 changes: 47 additions & 12 deletions packages/wasm/test/test.mjs
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,11 @@ const testBundle = async (t, bundle) => {
return func(t);
};

const setup = (t) => {
global.result = null;
global.t = t;
};

test('async compiling', async (t) => {
t.plan(2);

Expand All @@ -51,8 +56,7 @@ test('fetching WASM from separate file', async (t) => {
await bundle.write({ format: 'cjs', file: outputFile });
const glob = join(outputDir, `**/*.wasm`).split(sep).join(posix.sep);

global.result = null;
global.t = t;
setup(t);
await import(outputFile);
await global.result;
t.snapshot(await globby(glob));
Expand All @@ -61,6 +65,7 @@ test('fetching WASM from separate file', async (t) => {

test('complex module decoding', async (t) => {
t.plan(2);
setup(t);

const bundle = await rollup({
input: 'fixtures/complex.js',
Expand All @@ -71,6 +76,7 @@ test('complex module decoding', async (t) => {

test('sync compiling', async (t) => {
t.plan(2);
setup(t);

const bundle = await rollup({
input: 'fixtures/sync.js',
Expand All @@ -85,6 +91,7 @@ test('sync compiling', async (t) => {

test('imports', async (t) => {
t.plan(1);
setup(t);

const bundle = await rollup({
input: 'fixtures/imports.js',
Expand All @@ -99,6 +106,7 @@ test('imports', async (t) => {

test('worker', async (t) => {
t.plan(2);
setup(t);

const bundle = await rollup({
input: 'fixtures/worker.js',
Expand All @@ -120,6 +128,7 @@ test('worker', async (t) => {

test('injectHelper', async (t) => {
t.plan(4);
setup(t);

const injectImport = `import { _loadWasmModule } from ${JSON.stringify('\0wasmHelpers.js')};`;

Expand All @@ -146,12 +155,13 @@ test('injectHelper', async (t) => {
await testBundle(t, bundle);
});

test('target environment auto', async (t) => {
test('loader auto', async (t) => {
t.plan(5);
setup(t);

const bundle = await rollup({
input: 'fixtures/async.js',
plugins: [wasmPlugin({ targetEnv: 'auto' })]
plugins: [wasmPlugin({ loader: 'auto' })]
});
const code = await getCode(bundle);
await testBundle(t, bundle);
Expand All @@ -160,12 +170,13 @@ test('target environment auto', async (t) => {
t.true(code.includes(`fetch`));
});

test('target environment auto-inline', async (t) => {
test('loader auto-inline', async (t) => {
t.plan(6);
setup(t);

const bundle = await rollup({
input: 'fixtures/async.js',
plugins: [wasmPlugin({ targetEnv: 'auto-inline' })]
plugins: [wasmPlugin({ loader: 'auto-inline' })]
});
const code = await getCode(bundle);
await testBundle(t, bundle);
Expand All @@ -175,39 +186,61 @@ test('target environment auto-inline', async (t) => {
t.true(code.includes(`if (isNode)`));
});

test('target environment browser', async (t) => {
test('loader browser', async (t) => {
t.plan(4);
setup(t);

const bundle = await rollup({
input: 'fixtures/async.js',
plugins: [wasmPlugin({ targetEnv: 'browser' })]
plugins: [wasmPlugin({ loader: 'browser' })]
});
const code = await getCode(bundle);
await testBundle(t, bundle);
t.true(!code.includes(`require("`));
t.true(code.includes(`fetch`));
});

test('target environment node', async (t) => {
test('loader node', async (t) => {
t.plan(4);
setup(t);

const bundle = await rollup({
input: 'fixtures/async.js',
plugins: [wasmPlugin({ targetEnv: 'node' })]
plugins: [wasmPlugin({ loader: 'node' })]
});
const code = await getCode(bundle);
await testBundle(t, bundle);
t.true(code.includes(`require("`));
t.true(!code.includes(`fetch`));
});

test('loader custom', async (t) => {
t.plan(1);
setup(t);

function custom(sync, path, base64, imports) {
// eslint-disable-next-line no-console
console.log(`custom load: ${sync}, ${path}, ${base64}, ${imports}`);
}

const bundle = await rollup({
input: 'fixtures/async.js',
plugins: [wasmPlugin({ loader: custom })]
});

const code = await getCode(bundle);

t.true(code.includes('custom load'));
});

test('filename override', async (t) => {
t.plan(1);
setup(t);

const bundle = await rollup({
input: 'fixtures/async.js',
plugins: [
wasmPlugin({ maxFileSize: 0, targetEnv: 'node', fileName: 'start-[name]-suffix[extname]' })
wasmPlugin({ maxFileSize: 0, loader: 'node', fileName: 'start-[name]-suffix[extname]' })
]
});

Expand All @@ -220,6 +253,7 @@ test('filename override', async (t) => {

test('works as CJS plugin', async (t) => {
t.plan(2);
setup(t);
const require = createRequire(import.meta.url);
const wasmPluginCjs = require('current-package');
const bundle = await rollup({
Expand All @@ -233,10 +267,11 @@ test('works as CJS plugin', async (t) => {
if (!process.version.startsWith('v14')) {
test('avoid uncaught exception on file read', async (t) => {
t.plan(2);
setup(t);

const bundle = await rollup({
input: 'fixtures/complex.js',
plugins: [wasmPlugin({ maxFileSize: 0, targetEnv: 'node' })]
plugins: [wasmPlugin({ maxFileSize: 0, loader: 'node' })]
});

const raw = await getCode(bundle);
Expand Down
31 changes: 31 additions & 0 deletions packages/wasm/types/index.d.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,23 @@ import type { FilterPattern } from '@rollup/pluginutils';
*/
export type TargetEnv = 'auto' | 'auto-inline' | 'browser' | 'node';

/**
* The type for the plugin's loader function
*
* This is the function that ends up called when encountering a WASM import to load it and turn it into an usable object at runtime.
*
* @param {boolean} sync Whether the load should happen synchronously or not
* @param {string | null} filepath The path to the module.
* @param {string | null} src The base64-encoded source of the module
* @param {any} imports An object containing the module's imports
*/
export type WasmLoaderFunction = (
sync: boolean,
filepath: string | null,
src: string,
imports: any
) => void;

export interface RollupWasmOptions {
/**
* A picomatch pattern, or array of patterns, which specifies the files in the build the plugin
Expand Down Expand Up @@ -40,8 +57,22 @@ export interface RollupWasmOptions {
* A string which will be added in front of filenames when they are not inlined but are copied.
*/
publicPath?: string;
/**
* The loader used to process WASM modules.
*
* This plugin provides 4 default loaders:
* - `"auto"` will determine the environment at runtime and invoke the correct methods accordingly
* - `"auto-inline"` always inlines the Wasm and will decode it according to the environment
* - `"browser"` omits emitting code that requires node.js builtin modules that may play havoc on downstream bundlers
* - `"node"` omits emitting code that requires `fetch`
*
* Additionally, you can pass your own loader function if you need better control. The plugin expects a
* function with the following signature: `_loadWasmModule(sync: boolean, filepath: string, src: string, imports: any)`.
*/
loader?: TargetEnv | WasmLoaderFunction;
/**
* Configures what code is emitted to instantiate the Wasm (both inline and separate)
* @deprecated Use {@link RollupWasmOptions.loader}
*/
targetEnv?: TargetEnv;
}
Expand Down

0 comments on commit 226dfd6

Please sign in to comment.