Skip to content

Commit

Permalink
fix(cloudflare): wasm support with dynamic chunks (#1957)
Browse files Browse the repository at this point in the history
  • Loading branch information
pi0 authored Nov 28, 2023
1 parent 7ae0bcd commit 84c5418
Show file tree
Hide file tree
Showing 8 changed files with 172 additions and 74 deletions.
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -74,3 +74,4 @@ playground/firebase.json
test/fixture/functions

.pnpm-store
.wrangler
8 changes: 5 additions & 3 deletions package.json
Original file line number Diff line number Diff line change
Expand Up @@ -45,12 +45,12 @@
"dev:start": "node playground/.output/server/index.mjs",
"lint": "eslint --cache --ext .ts,.mjs,.cjs . && prettier -c src test",
"lint:fix": "eslint --cache --fix --ext .ts,.mjs,.cjs . && prettier --write -c src test",
"nitro": "NODE_OPTIONS=\"--enable-source-maps\" jiti ./src/cli/index.ts",
"nitro": "JITI_ESM_RESOLVE=1 NODE_OPTIONS=\"--enable-source-maps\" jiti ./src/cli/index.ts",
"prepack": "pnpm build",
"release": "pnpm test && pnpm build && changelogen --release && pnpm publish && git push --follow-tags",
"stub": "unbuild --stub",
"test": "pnpm lint && pnpm vitest-es run --silent",
"test:fixture:types": "pnpm stub && jiti ./test/scripts/gen-fixture-types.ts && cd test/fixture && tsc --noEmit",
"test:fixture:types": "pnpm stub && JITI_ESM_RESOLVE=1 jiti ./test/scripts/gen-fixture-types.ts && cd test/fixture && tsc --noEmit",
"test:types": "tsc --noEmit && pnpm test:fixture:types",
"vitest-es": "NODE_OPTIONS=\"--enable-source-maps --experimental-vm-modules\" vitest"
},
Expand Down Expand Up @@ -84,6 +84,7 @@
"dot-prop": "^8.0.2",
"esbuild": "^0.19.6",
"escape-string-regexp": "^5.0.0",
"estree-walker": "^3.0.3",
"etag": "^1.8.1",
"fs-extra": "^11.1.1",
"globby": "^14.0.0",
Expand Down Expand Up @@ -127,6 +128,7 @@
"@azure/functions": "^3.5.1",
"@cloudflare/workers-types": "^4.20231025.0",
"@types/aws-lambda": "^8.10.126",
"@types/estree": "^1.0.5",
"@types/etag": "^1.8.3",
"@types/fs-extra": "^11.0.4",
"@types/node-fetch": "^2.6.9",
Expand Down Expand Up @@ -172,4 +174,4 @@
]
}
}
}
}
6 changes: 6 additions & 0 deletions pnpm-lock.yaml

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

5 changes: 1 addition & 4 deletions src/rollup/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -361,10 +361,7 @@ export const plugins = [
return { id: _resolved, external: false };
}
}
if (
!resolved ||
(resolved.external && resolved.resolvedBy !== "nitro:wasm-import")
) {
if (!resolved || (resolved.external && !id.endsWith(".wasm"))) {
throw new Error(
`Cannot resolve ${JSON.stringify(id)} from ${JSON.stringify(
from
Expand Down
185 changes: 137 additions & 48 deletions src/rollup/plugins/wasm.ts
Original file line number Diff line number Diff line change
@@ -1,71 +1,160 @@
import { createHash } from "node:crypto";
import { extname, basename } from "node:path";
import { promises as fs } from "node:fs";
import { promises as fs, existsSync } from "node:fs";
import { basename, normalize } from "pathe";
import type { Plugin } from "rollup";
import wasmBundle from "@rollup/plugin-wasm";
import { isWindows } from "std-env";
import MagicString from "magic-string";
import { walk } from "estree-walker";
import { WasmOptions } from "../../types";

const PLUGIN_NAME = "nitro:wasm-import";
const wasmRegex = /\.wasm$/;

export function wasm(options: WasmOptions): Plugin {
return options.esmImport && !isWindows /* TODO */
? wasmImport()
: wasmBundle(options.rollup);
return options.esmImport ? wasmImport() : wasmBundle(options.rollup);
}

const WASM_IMPORT_PREFIX = "\0nitro:wasm/";

export function wasmImport(): Plugin {
const copies = Object.create(null);
type WasmAssetInfo = {
fileName: string;
id: string;
source: Buffer;
hash: string;
};

return {
name: PLUGIN_NAME,
async resolveId(id: string, importer: string) {
if (copies[id]) {
return {
id: copies[id].publicFilepath,
external: true,
const wasmSources = new Map<string /* sourceFile */, WasmAssetInfo>();
const wasmImports = new Map<string /* id */, WasmAssetInfo>();
const wasmGlobals = new Map<string /* global id */, WasmAssetInfo>();

return <Plugin>{
name: "nitro:wasm",
async resolveId(id, importer, options) {
// Only handle .wasm imports
if (!id.endsWith(".wasm")) {
return null;
}

// Resolve the source file real path
const sourceFile = await this.resolve(id, importer, options).then((r) =>
r?.id ? normalize(r.id) : null
);
if (!sourceFile || !existsSync(sourceFile)) {
return null;
}

// Read (cached) Asset
let wasmAsset: WasmAssetInfo | undefined = wasmSources.get(sourceFile);
if (!wasmAsset) {
wasmAsset = {
id: WASM_IMPORT_PREFIX + sourceFile,
fileName: "",
source: undefined,
hash: "",
};
wasmSources.set(sourceFile, wasmAsset);
wasmImports.set(wasmAsset.id, wasmAsset);

wasmAsset.source = await fs.readFile(sourceFile);
wasmAsset.hash = sha1(wasmAsset.source);
const _baseName = basename(sourceFile, ".wasm");
wasmAsset.fileName = `wasm/${_baseName}-${wasmAsset.hash}.wasm`;

await this.emitFile({
type: "asset",
source: wasmAsset.source,
fileName: wasmAsset.fileName,
});
}

// Resolve as external
return {
id: wasmAsset.id,
external: true,
};
},
renderChunk(code, chunk) {
if (!code.includes(WASM_IMPORT_PREFIX)) {
return null;
}
if (wasmRegex.test(id)) {
const { id: filepath } = (await this.resolve(id, importer)) || {};
if (!filepath || filepath === id) {

const s = new MagicString(code);

const resolveImport = (specifier?: string) => {
if (
typeof specifier !== "string" ||
!specifier.startsWith(WASM_IMPORT_PREFIX)
) {
return null;
}
const buffer = await fs.readFile(filepath);
const hash = createHash("sha1")
.update(buffer)
.digest("hex")
.slice(0, 16);
const ext = extname(filepath);
const name = basename(filepath, ext);

const outputFileName = `wasm/${name}-${hash}${ext}`;
const publicFilepath = `./${outputFileName}`;
const asset = wasmImports.get(specifier);
if (!asset) {
return null;
}
const nestedLevel = chunk.fileName.split("/").length - 1;
return (
(nestedLevel ? "../".repeat(nestedLevel) : "./") + asset.fileName
);
};

copies[id] = {
filename: outputFileName,
publicFilepath,
buffer,
};
walk(this.parse(code) as any, {
enter(node, parent, prop, index) {
if (
// prettier-ignore
(node.type === "ImportDeclaration" || node.type === "ImportExpression") &&
"value" in node.source && typeof node.source.value === "string" &&
"start" in node.source && typeof node.source.start === "number" &&
"end" in node.source && typeof node.source.end === "number"
) {
const resolved = resolveImport(node.source.value);
if (resolved) {
// prettier-ignore
s.update(node.source.start, node.source.end, JSON.stringify(resolved));
}
}
},
});

if (s.hasChanged()) {
return {
id: publicFilepath,
external: true,
code: s.toString(),
map: s.generateMap({ includeContent: true }),
};
}
},
async generateBundle() {
await Promise.all(
Object.keys(copies).map(async (name) => {
const copy = copies[name];
await this.emitFile({
type: "asset",
source: copy.buffer,
fileName: copy.filename,
});
})
);
// --- [temporary] IIFE/UMD support for cloudflare (non module targets) ---
renderStart(options) {
if (options.format === "iife" || options.format === "umd") {
for (const [importName, wasmAsset] of wasmImports.entries()) {
if (!(importName in options.globals)) {
const globalName = `_wasm_${wasmAsset.hash}`;
wasmGlobals.set(globalName, wasmAsset);
options.globals[importName] = globalName;
}
}
}
},
generateBundle(options, bundle) {
if (wasmGlobals.size > 0) {
for (const [fileName, chunkInfo] of Object.entries(bundle)) {
if (chunkInfo.type !== "chunk" || !chunkInfo.isEntry) {
continue;
}
const imports: string[] = [];
for (const [globalName, wasmAsset] of wasmGlobals.entries()) {
if (chunkInfo.code.includes(globalName)) {
imports.push(
`import ${globalName} from "${wasmAsset.fileName}";`
);
}
}
if (imports.length > 0) {
chunkInfo.code = imports.join("\n") + "\n" + chunkInfo.code;
}
}
}
},
};
}

function sha1(source: Buffer) {
return createHash("sha1").update(source).digest("hex").slice(0, 16);
}
17 changes: 17 additions & 0 deletions test/fixture/routes/wasm/dynamic.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
export default defineLazyEventHandler(async () => {
const { sum } = await importWasm(import("~/wasm/sum.wasm" as string));
return eventHandler(() => {
return `2+3=${sum(2, 3)}`;
});
});

// TODO: Extract as reusable utility once stable
async function importWasm(input: any) {
const _input = await input;
const _module = _input.default || _input;
const _instance =
typeof _module === "function"
? await _module({}).then((r) => r.instance || r)
: await WebAssembly.instantiate(_module, {});
return _instance.exports;
}
15 changes: 0 additions & 15 deletions test/fixture/routes/wasm/sum.ts

This file was deleted.

9 changes: 5 additions & 4 deletions test/tests.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@ import { expect, it, afterAll, beforeAll, describe } from "vitest";
import { fileURLToPath } from "mlly";
import { joinURL } from "ufo";
import { defu } from "defu";
import { isWindows } from "std-env";
import * as _nitro from "../src";
import type { Nitro } from "../src";

Expand Down Expand Up @@ -632,10 +631,12 @@ export function testNitro(
});

describe("wasm", () => {
it.skipIf(ctx.isWorker || ctx.preset === "deno-server" || isWindows)(
"sum works",
it.skipIf(ctx.isWorker || ctx.preset === "deno-server")(
"dynamic import wasm",
async () => {
expect((await callHandler({ url: "/wasm/sum" })).data).toBe("2+3=5");
expect((await callHandler({ url: "/wasm/dynamic" })).data).toBe(
"2+3=5"
);
}
);
});
Expand Down

0 comments on commit 84c5418

Please sign in to comment.