feat: support WebAssembly (Wasm) imports in ESM modules

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ### Features
 
+- `[jest-runtime]` Support 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))

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 Wasm test with native ESM 1`] = `
+"Test Suites: 1 passed, 1 total
+Tests:       4 passed, 4 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 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,37 @@
+/**
+ * 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.
+ */
+
+// the point here is that it's the node core module
+// eslint-disable-next-line no-restricted-imports
+import {readFileSync} from 'fs';
+// The file was generated by wasm-pack
+import {getAnswer} from '../42.wasm';
+
+const wasmFileBuffer = readFileSync('42.wasm');
+
+test('supports native wasm imports', () => {
+  expect(getAnswer()).toBe(42);
+});
+
+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.getAnswer()).toBe(42);
+});
+
+test('imports from "data:text/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:text/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

@@ -398,9 +398,12 @@ 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,
+    return (
+      path.endsWith('.wasm') ||
+      Resolver.unstable_shouldLoadAsEsm(
+        path,
+        this._config.extensionsToTreatAsEsm,
+      )
     );
   }
 
@@ -441,6 +444,19 @@ export default class Runtime {
         'Promise initialization should be sync - please report this bug to Jest!',
       );
 
+      if (modulePath.endsWith('.wasm')) {
+        const wasm = this._importWasmModule(
+          fs.readFileSync(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 +583,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);
@@ -1640,6 +1667,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.linkAndEvaluateModule(
+          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;