feat: support WebAssembly (Wasm) imports in ESM modules (#13505)

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ### Features
 
+- `[jest-runtime]` Support WebAssembly (Wasm) imports in ESM modules ([#13505](https://github.com/facebook/jest/pull/13505))
+
 ### Fixes
 
 - `[jest-mock]` Treat cjs modules as objects so they can be mocked ([#13513](https://github.com/facebook/jest/pull/13513))

docs/ECMAScriptModules.md

@@ -22,6 +22,8 @@ With the warnings out of the way, this is how you activate ESM support in your t
 
    If you use Yarn, you can use `yarn node --experimental-vm-modules $(yarn bin jest)`. This command will also work if you use [Yarn Plug'n'Play](https://yarnpkg.com/features/pnp).
 
+   If your codebase includes ESM imports from `*.wasm` files, you do _not_ need to pass `--experimental-wasm-modules` to `node`. Current implementation of WebAssembly imports in Jest relies on experimental VM modules, however, this may change in the future.
+
 1. Beyond that, we attempt to follow `node`'s logic for activating "ESM mode" (such as looking at `type` in `package.json` or `.mjs` files), see [their docs](https://nodejs.org/api/esm.html#esm_enabling) for details.
 1. If you want to treat other file extensions (such as `.jsx` or `.ts`) as ESM, please use the [`extensionsToTreatAsEsm` option](Configuration.md#extensionstotreatasesm-arraystring).
 

e2e/__tests__/__snapshots__/nativeEsm.test.ts.snap

@@ -8,9 +8,17 @@ Time:        <<REPLACED>>
 Ran all test suites matching /native-esm-deep-cjs-reexport.test.js/i."
 `;
 
+exports[`runs WebAssembly (Wasm) test with native ESM 1`] = `
+"Test Suites: 1 passed, 1 total
+Tests:       5 passed, 5 total
+Snapshots:   0 total
+Time:        <<REPLACED>>
+Ran all test suites matching /native-esm-wasm.test.js/i."
+`;
+
 exports[`runs test with native ESM 1`] = `
 "Test Suites: 1 passed, 1 total
-Tests:       34 passed, 34 total
+Tests:       33 passed, 33 total
 Snapshots:   0 total
 Time:        <<REPLACED>>
 Ran all test suites matching /native-esm.test.js/i."

e2e/__tests__/nativeEsm.test.ts

@@ -67,3 +67,15 @@ onNodeVersions('>=16.9.0', () => {
     expect(exitCode).toBe(0);
   });
 });
+
+test('runs WebAssembly (Wasm) test with native ESM', () => {
+  const {exitCode, stderr, stdout} = runJest(DIR, ['native-esm-wasm.test.js'], {
+    nodeOptions: '--experimental-vm-modules --no-warnings',
+  });
+
+  const {summary} = extractSummary(stderr);
+
+  expect(summary).toMatchSnapshot();
+  expect(stdout).toBe('');
+  expect(exitCode).toBe(0);
+});

e2e/native-esm/__tests__/native-esm-wasm.test.js

@@ -0,0 +1,56 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import {readFileSync} from 'node:fs';
+// file origin: https://github.com/mdn/webassembly-examples/blob/2f2163287f86fe29deb162335bccca7d5d95ca4f/understanding-text-format/add.wasm
+// source code: https://github.com/mdn/webassembly-examples/blob/2f2163287f86fe29deb162335bccca7d5d95ca4f/understanding-text-format/add.was
+import {add} from '../add.wasm';
+
+const wasmFileBuffer = readFileSync('add.wasm');
+
+test('supports native wasm imports', () => {
+  expect(add(1, 2)).toBe(3);
+
+  // because arguments are i32 (signed), fractional part is truncated
+  expect(add(0.99, 1.01)).toBe(1);
+
+  // because return value is i32 (signed), (2^31 - 1) + 1 overflows and becomes -2^31
+  expect(add(Math.pow(2, 31) - 1, 1)).toBe(-Math.pow(2, 31));
+
+  // invalid or missing arguments are treated as 0
+  expect(add('hello', 'world')).toBe(0);
+  expect(add()).toBe(0);
+  expect(add(null)).toBe(0);
+  expect(add({}, [])).toBe(0);
+
+  // redundant arguments are silently ignored
+  expect(add(1, 2, 3)).toBe(3);
+});
+
+test('supports dynamic wasm imports', async () => {
+  const {add: dynamicAdd} = await import('../add.wasm');
+  expect(dynamicAdd(1, 2)).toBe(3);
+});
+
+test('supports imports from "data:application/wasm" URI with base64 encoding', async () => {
+  const importedWasmModule = await import(
+    `data:application/wasm;base64,${wasmFileBuffer.toString('base64')}`
+  );
+  expect(importedWasmModule.add(0, 42)).toBe(42);
+});
+
+test('imports from "data:application/wasm" URI without explicit encoding fail', async () => {
+  await expect(() =>
+    import(`data:application/wasm,${wasmFileBuffer.toString('base64')}`),
+  ).rejects.toThrow('Missing data URI encoding');
+});
+
+test('imports from "data:application/wasm" URI with invalid encoding fail', async () => {
+  await expect(() =>
+    import('data:application/wasm;charset=utf-8,oops'),
+  ).rejects.toThrow('Invalid data URI encoding: charset=utf-8');
+});

e2e/native-esm/__tests__/native-esm.test.js

@@ -255,12 +255,6 @@ test('imports from "data:text/javascript" URI with invalid data fail', async ()
   ).rejects.toThrow("Unexpected token '.'");
 });
 
-test('imports from "data:application/wasm" URI not supported', async () => {
-  await expect(() =>
-    import('data:application/wasm,96cafe00babe'),
-  ).rejects.toThrow('WASM is currently not supported');
-});
-
 test('supports imports from "data:application/json" URI', async () => {
   const data = await import('data:application/json,{"foo": "bar"}');
   expect(data.default).toEqual({foo: 'bar'});

packages/jest-runtime/src/index.ts

@@ -134,6 +134,8 @@ const getModuleNameMapper = (config: Config.ProjectConfig) => {
   return null;
 };
 
+const isWasm = (modulePath: string): boolean => modulePath.endsWith('.wasm');
+
 const unmockRegExpCache = new WeakMap();
 
 const EVAL_RESULT_VARIABLE = 'Object.<anonymous>';
@@ -154,6 +156,7 @@ const supportsNodeColonModulePrefixInRequire = (() => {
 
 export default class Runtime {
   private readonly _cacheFS: Map<string, string>;
+  private readonly _cacheFSBuffer = new Map<string, Buffer>();
   private readonly _config: Config.ProjectConfig;
   private readonly _globalConfig?: Config.GlobalConfig;
   private readonly _coverageOptions: ShouldInstrumentOptions;
@@ -397,10 +400,13 @@ export default class Runtime {
   }
 
   // unstable as it should be replaced by https://github.com/nodejs/modules/issues/393, and we don't want people to use it
-  unstable_shouldLoadAsEsm(path: string): boolean {
-    return Resolver.unstable_shouldLoadAsEsm(
-      path,
-      this._config.extensionsToTreatAsEsm,
+  unstable_shouldLoadAsEsm(modulePath: string): boolean {
+    return (
+      isWasm(modulePath) ||
+      Resolver.unstable_shouldLoadAsEsm(
+        modulePath,
+        this._config.extensionsToTreatAsEsm,
+      )
     );
   }
 
@@ -441,6 +447,19 @@ export default class Runtime {
         'Promise initialization should be sync - please report this bug to Jest!',
       );
 
+      if (isWasm(modulePath)) {
+        const wasm = this._importWasmModule(
+          this.readFileBuffer(modulePath),
+          modulePath,
+          context,
+        );
+
+        this._esmoduleRegistry.set(cacheKey, wasm);
+
+        transformResolve();
+        return wasm;
+      }
+
       if (this._resolver.isCoreModule(modulePath)) {
         const core = this._importCoreModule(modulePath, context);
         this._esmoduleRegistry.set(cacheKey, core);
@@ -567,56 +586,67 @@ export default class Runtime {
       }
 
       const mime = match.groups.mime;
-      if (mime === 'application/wasm') {
-        throw new Error('WASM is currently not supported');
-      }
-
       const encoding = match.groups.encoding;
-      let code = match.groups.code;
-      if (!encoding || encoding === 'charset=utf-8') {
-        code = decodeURIComponent(code);
-      } else if (encoding === 'base64') {
-        code = Buffer.from(code, 'base64').toString();
-      } else {
-        throw new Error(`Invalid data URI encoding: ${encoding}`);
-      }
-
       let module;
-      if (mime === 'application/json') {
-        module = new SyntheticModule(
-          ['default'],
-          function () {
-            const obj = JSON.parse(code);
-            // @ts-expect-error: TS doesn't know what `this` is
-            this.setExport('default', obj);
-          },
-          {context, identifier: specifier},
+
+      if (mime === 'application/wasm') {
+        if (!encoding) {
+          throw new Error('Missing data URI encoding');
+        }
+        if (encoding !== 'base64') {
+          throw new Error(`Invalid data URI encoding: ${encoding}`);
+        }
+        module = await this._importWasmModule(
+          Buffer.from(match.groups.code, 'base64'),
+          specifier,
+          context,
         );
       } else {
-        module = new SourceTextModule(code, {
-          context,
-          identifier: specifier,
-          importModuleDynamically: async (
-            specifier: string,
-            referencingModule: VMModule,
-          ) => {
-            invariant(
-              runtimeSupportsVmModules,
-              'You need to run with a version of node that supports ES Modules in the VM API. See https://jestjs.io/docs/ecmascript-modules',
-            );
-            const module = await this.resolveModule(
-              specifier,
-              referencingModule.identifier,
-              referencingModule.context,
-            );
+        let code = match.groups.code;
+        if (!encoding || encoding === 'charset=utf-8') {
+          code = decodeURIComponent(code);
+        } else if (encoding === 'base64') {
+          code = Buffer.from(code, 'base64').toString();
+        } else {
+          throw new Error(`Invalid data URI encoding: ${encoding}`);
+        }
 
-            return this.linkAndEvaluateModule(module);
-          },
-          initializeImportMeta(meta: ImportMeta) {
-            // no `jest` here as it's not loaded in a file
-            meta.url = specifier;
-          },
-        });
+        if (mime === 'application/json') {
+          module = new SyntheticModule(
+            ['default'],
+            function () {
+              const obj = JSON.parse(code);
+              // @ts-expect-error: TS doesn't know what `this` is
+              this.setExport('default', obj);
+            },
+            {context, identifier: specifier},
+          );
+        } else {
+          module = new SourceTextModule(code, {
+            context,
+            identifier: specifier,
+            importModuleDynamically: async (
+              specifier: string,
+              referencingModule: VMModule,
+            ) => {
+              invariant(
+                runtimeSupportsVmModules,
+                'You need to run with a version of node that supports ES Modules in the VM API. See https://jestjs.io/docs/ecmascript-modules',
+              );
+              const module = await this.resolveModule(
+                specifier,
+                referencingModule.identifier,
+                referencingModule.context,
+              );
+
+              return this.linkAndEvaluateModule(module);
+            },
+            initializeImportMeta(meta: ImportMeta) {
+              // no `jest` here as it's not loaded in a file
+              meta.url = specifier;
+            },
+          });
+        }
       }
 
       this._esmoduleRegistry.set(specifier, module);
@@ -1117,6 +1147,7 @@ export default class Runtime {
     this._cjsNamedExports.clear();
     this._moduleMockRegistry.clear();
     this._cacheFS.clear();
+    this._cacheFSBuffer.clear();
 
     if (
       this._coverageOptions.collectCoverage &&
@@ -1640,6 +1671,50 @@ export default class Runtime {
     return evaluateSyntheticModule(module);
   }
 
+  private async _importWasmModule(
+    source: Buffer,
+    identifier: string,
+    context: VMContext,
+  ) {
+    const wasmModule = await WebAssembly.compile(source);
+
+    const exports = WebAssembly.Module.exports(wasmModule);
+    const imports = WebAssembly.Module.imports(wasmModule);
+
+    const moduleLookup: Record<string, VMModule> = {};
+    for (const {module} of imports) {
+      if (moduleLookup[module] === undefined) {
+        moduleLookup[module] = await this.loadEsmModule(
+          await this.resolveModule(module, identifier, context),
+        );
+      }
+    }
+
+    const syntheticModule = new SyntheticModule(
+      exports.map(({name}) => name),
+      function () {
+        const importsObject: WebAssembly.Imports = {};
+        for (const {module, name} of imports) {
+          if (!importsObject[module]) {
+            importsObject[module] = {};
+          }
+          importsObject[module][name] = moduleLookup[module].namespace[name];
+        }
+        const wasmInstance = new WebAssembly.Instance(
+          wasmModule,
+          importsObject,
+        );
+        for (const {name} of exports) {
+          // @ts-expect-error: TS doesn't know what `this` is
+          this.setExport(name, wasmInstance.exports[name]);
+        }
+      },
+      {context, identifier},
+    );
+
+    return syntheticModule;
+  }
+
   private _getMockedNativeModule(): typeof nativeModule.Module {
     if (this._moduleImplementation) {
       return this._moduleImplementation;
@@ -2305,11 +2380,24 @@ export default class Runtime {
     };
   }
 
+  private readFileBuffer(filename: string): Buffer {
+    let source = this._cacheFSBuffer.get(filename);
+
+    if (!source) {
+      source = fs.readFileSync(filename);
+
+      this._cacheFSBuffer.set(filename, source);
+    }
+
+    return source;
+  }
+
   private readFile(filename: string): string {
     let source = this._cacheFS.get(filename);
 
     if (!source) {
-      source = fs.readFileSync(filename, 'utf8');
+      const buffer = this.readFileBuffer(filename);
+      source = buffer.toString('utf8');
 
       this._cacheFS.set(filename, source);
     }