diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 91615a2..709f514 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -18,6 +18,8 @@ jobs: - uses: actions/checkout@v5 - run: npm i -g --force corepack - run: corepack enable + - run: sudo apt install -y binaryen + if: ${{ matrix.os == 'ubuntu-latest' }} - uses: actions/setup-node@v4 with: node-version: 22 diff --git a/examples/build.mjs b/examples/build.mjs index 77278bd..c88236d 100644 --- a/examples/build.mjs +++ b/examples/build.mjs @@ -1,5 +1,5 @@ import { fileURLToPath } from "node:url"; -import fs from 'node:fs/promises'; +import fs from "node:fs/promises"; import Module from "node:module"; import { main as asc } from "assemblyscript/asc"; @@ -22,7 +22,7 @@ async function compile(name) { async function compileWat(name) { const module = wabt.parseWat(`${name}.wat`, await fs.readFile(`${name}.wat`)); module.resolveNames(); - const binaryOutput = module.toBinary({write_debug_names:true}); + const binaryOutput = module.toBinary({ write_debug_names: true }); const binaryBuffer = binaryOutput.buffer; await fs.writeFile(`${name}.wasm`, binaryBuffer); } @@ -32,3 +32,4 @@ process.chdir(fileURLToPath(new URL(".", import.meta.url))); await compile("sum"); await compile("rand"); await compileWat("add-esmi"); +await compileWat("treeshake"); diff --git a/examples/treeshake.wasm b/examples/treeshake.wasm new file mode 100644 index 0000000..3f202ac Binary files /dev/null and b/examples/treeshake.wasm differ diff --git a/examples/treeshake.wat b/examples/treeshake.wat new file mode 100644 index 0000000..ca290f5 --- /dev/null +++ b/examples/treeshake.wat @@ -0,0 +1,8 @@ +(module + (func (export "functionOne") (result i32) + i32.const 42 + ) + (func (export "functionTwo") (result i32) + i32.const 0xDEADBEEF + ) +) diff --git a/src/plugin/index.ts b/src/plugin/index.ts index e43ffef..294ec79 100644 --- a/src/plugin/index.ts +++ b/src/plugin/index.ts @@ -5,6 +5,7 @@ import type { RenderedChunk, Plugin as RollupPlugin } from "rollup"; import { createUnplugin } from "unplugin"; import { parseWasm } from "../tools"; import { getWasmESMBinding, getWasmModuleBinding } from "./runtime/binding"; +import { treeshake, filterExports } from "./runtime/treeshake"; import { getPluginUtils } from "./runtime/utils"; import { sha1, @@ -139,8 +140,30 @@ const unplugin = createUnplugin((opts) => { }; }, renderChunk(code: string, chunk: RenderedChunk) { - if (!opts.esmImport) { - return; + if (chunk.type === "chunk" && opts.treeshake) { + for (const [id, module] of Object.entries(chunk.modules)) { + if (!WASM_ID_RE.test(id)) { + continue; + } + // Find the asset and tree shake it. + let assetKey; + for (const [key, value] of Object.entries(assets)) { + if (value.id === id) { + assetKey = key; + break; + } + } + if (!assetKey) { + throw new Error("Could not find asset."); + } + const asset = assets[assetKey]; + const buffer = treeshake( + asset.source, + filterExports(asset.exports, module.removedExports), + opts, + ); + asset.source = buffer; + } } if ( @@ -160,10 +183,16 @@ const unplugin = createUnplugin((opts) => { if (!asset) { return; } - const nestedLevel = - chunk.fileName.split("/").filter(Boolean /* handle // */).length - 1; - const relativeId = - (nestedLevel ? "../".repeat(nestedLevel) : "./") + asset.name; + let relativeId; + if (opts.esmImport) { + const nestedLevel = + chunk.fileName.split("/").filter(Boolean /* handle // */).length - + 1; + relativeId = + (nestedLevel ? "../".repeat(nestedLevel) : "./") + asset.name; + } else { + relativeId = asset.source.toString("base64"); + } return { relativeId, asset, diff --git a/src/plugin/runtime/binding.ts b/src/plugin/runtime/binding.ts index 3b15034..8b80cd8 100644 --- a/src/plugin/runtime/binding.ts +++ b/src/plugin/runtime/binding.ts @@ -34,7 +34,7 @@ export default _mod; ` : /* js */ ` import { base64ToUint8Array } from "${UMWASM_HELPERS_ID}"; -const _data = base64ToUint8Array("${asset.source.toString("base64")}"); +const _data = base64ToUint8Array("${UNWASM_EXTERNAL_PREFIX}${asset.name}"); const _mod = new WebAssembly.Module(_data); export default _mod; `; @@ -60,7 +60,7 @@ import { base64ToUint8Array } from "${UMWASM_HELPERS_ID}"; ${importsCode} function _instantiate(imports = _imports) { - const _data = base64ToUint8Array("${asset.source.toString("base64")}") + const _data = base64ToUint8Array("${UNWASM_EXTERNAL_PREFIX}${asset.name}") return WebAssembly.instantiate(_data, imports) } `; } diff --git a/src/plugin/runtime/treeshake.ts b/src/plugin/runtime/treeshake.ts new file mode 100644 index 0000000..be5135c --- /dev/null +++ b/src/plugin/runtime/treeshake.ts @@ -0,0 +1,79 @@ +import fs, { mkdirSync } from "node:fs"; +import { execSync } from "node:child_process"; +import { UnwasmPluginOptions } from "../shared"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; + +export function treeshake( + source: Buffer, + exports: string[], + opts: UnwasmPluginOptions, +): Buffer { + // Create the wasm-metadce graph. + const graph = [ + { + name: "outside", + root: true, + reaches: [...exports], + }, + ] as { name: string; root?: boolean; reaches?: string[]; export?: string }[]; + for (const exportName of exports) { + graph.push({ + name: exportName, + export: exportName, + }); + } + + // Wasm Meta DCE binary + const wasmMetaDCEBin = opts.commands?.wasmMetaDCE || "wasm-metadce"; + + // Temporary files + const tmpDir = join( + tmpdir(), + "unwasm-" + Math.random().toString(36).slice(2), + ); + const graphFile = join(tmpDir, "graph.json"); + const wasmFile = join(tmpDir, "input.wasm"); + const outputFile = join(tmpDir, "output.wasm"); + mkdirSync(tmpDir, { recursive: true }); + fs.writeFileSync(graphFile, JSON.stringify(graph)); + fs.writeFileSync(wasmFile, source); + + // Execute the wasm-metadce command + try { + execSync( + `${wasmMetaDCEBin} ${wasmFile} --graph-file ${graphFile} -o ${outputFile}`, + { stdio: "ignore" }, + ); + return fs.readFileSync(outputFile); + } catch (error) { + throw new Error( + [ + error instanceof Error ? error.message : String(error), + `Hint: Make sure "wasm-metadce" (part of binaryen) is installed and on your PATH,`, + `or set commands.wasmMetaDCE option to the full path of the executable.`, + ].join(" "), + { cause: error }, + ); + } finally { + fs.rmSync(tmpDir, { force: true, recursive: true }); + } +} + +/** + * Removes elements from one array that are present in another array. + * + * @param exports - The array of strings to filter. + * @param removedExports - The array of strings to remove from the exports array. + * @returns A new array containing elements from exports with removedExports removed. + */ +export function filterExports(exports: string[], removedExports: string[]) { + // Create a Set from removedExports for efficient lookup. + const removedExportsSet = new Set(removedExports); + const filteredExports = exports.filter((exportItem) => { + // Keep the item if it's NOT in the removedExportsSet. + return !removedExportsSet.has(exportItem); + }); + + return filteredExports; +} diff --git a/src/plugin/shared.ts b/src/plugin/shared.ts index 53fa69e..1c876f4 100644 --- a/src/plugin/shared.ts +++ b/src/plugin/shared.ts @@ -16,6 +16,18 @@ export interface UnwasmPluginOptions { * @default false */ lazy?: boolean; + + /** + * Enable treeshaking of Wasm files using Binaryen's wasm-metadce. + * + * @default false + */ + treeshake?: boolean; + + commands?: { + /** Path to wasm-metadce binary. Defaults to `wasm-metadce` from the system. */ + wasmMetaDCE?: string; + }; } export type WasmAsset = { diff --git a/test/fixture/treeshake.mjs b/test/fixture/treeshake.mjs new file mode 100644 index 0000000..7202e64 --- /dev/null +++ b/test/fixture/treeshake.mjs @@ -0,0 +1,8 @@ +import { functionOne } from "@fixture/wasm/treeshake.wasm"; + +export function test() { + if (functionOne(1) !== 42) { + return "FAILED: call export"; + } + return "OK"; +} diff --git a/test/node_modules/@fixture/wasm/treeshake.wasm b/test/node_modules/@fixture/wasm/treeshake.wasm new file mode 100644 index 0000000..3f202ac Binary files /dev/null and b/test/node_modules/@fixture/wasm/treeshake.wasm differ diff --git a/test/plugin.test.ts b/test/plugin.test.ts index 715ce49..2938e45 100644 --- a/test/plugin.test.ts +++ b/test/plugin.test.ts @@ -6,6 +6,7 @@ import { evalModule } from "mlly"; import { nodeResolve as rollupNodeResolve } from "@rollup/plugin-node-resolve"; import { rollup } from "rollup"; import { UnwasmPluginOptions, rollup as unwasmRollup } from "../src/plugin"; +import { execSync } from "node:child_process"; const r = (p: string) => fileURLToPath(new URL(p, import.meta.url)); @@ -75,6 +76,39 @@ describe("plugin:rollup", () => { }), ); }); + + // const hasBin = + let wasmMetaDCEBinExists: boolean; + try { + execSync("wasm-metadce --version"); + wasmMetaDCEBinExists = true; + } catch { + wasmMetaDCEBinExists = false; + } + + it.skipIf(!wasmMetaDCEBinExists)("treeshake", async () => { + const { output } = await _rollupBuild( + "fixture/treeshake.mjs", + "rollup-inline", + { + treeshake: true, + }, + ); + const code = output[0].code; + // Check that the unused export 'functionTwo' is removed from the Wasm + // module. + const wasm = getBase64WasmModule(code); + expect(wasm).toBeTruthy(); + const module = await WebAssembly.compile(Buffer.from(wasm, "base64")); + const exports = WebAssembly.Module.exports(module); + expect(exports.length).toEqual(1); + expect(exports[0].name).toEqual("functionOne"); + // Ensure that it still runs after dce. + const mod = await evalModule(code, { + url: r("fixture/treeshake.mjs"), + }); + expect(mod.test()).toBe("OK"); + }); }); // --- Utils --- @@ -115,3 +149,12 @@ export default { await mf.dispose(); return res; } + +function getBase64WasmModule(code: string) { + const start = code.indexOf('base64ToUint8Array("'); + if (start === -1) { + return false; + } + const end = code.indexOf('")', start); + return code.slice(start + 20, end); +}