feat: support treeshaking wasm modules

.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

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");

examples/treeshake.wat

@@ -0,0 +1,8 @@
+(module
+  (func (export "functionOne") (result i32)
+    i32.const 42
+  )
+  (func (export "functionTwo") (result i32)
+    i32.const 0xDEADBEEF
+  )
+)

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<UnwasmPluginOptions>((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<UnwasmPluginOptions>((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,

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)  }
   `;
 }

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;
+}

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 = {

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";
+}

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);
+}