build_tools/wasm/wasm_test_main.mjs

// Copyright 2026 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

// Generic test entry point for gtest/gbenchmark wasm binaries under WASI.
//
// The bundler appends this file after inlined companion code and the generated
// createWasmImports() function, so createWasmImports is available in module
// scope. The bundler also defines __IREE_WASM_BINARY with the .wasm filename.
//
// Provides WASI preview1 syscalls via node:wasi, merges companion imports on
// top, instantiates the wasm binary, and propagates the exit code.

import {readFileSync} from 'node:fs';
import {dirname, resolve} from 'node:path';
import {fileURLToPath} from 'node:url';
import {WASI} from 'node:wasi';

const scriptDirectory = dirname(fileURLToPath(import.meta.url));
const wasmPath = resolve(scriptDirectory, __IREE_WASM_BINARY);

// Build preopens for directories the test needs to access. Bazel sets
// TEST_TMPDIR for scratch space and XML_OUTPUT_FILE for gtest XML output.
// Without preopening these, WASI's sandbox blocks stat/open calls and gtest
// crashes during teardown.
const preopens = {};
if (process.env.TEST_TMPDIR) {
  preopens[process.env.TEST_TMPDIR] = process.env.TEST_TMPDIR;
}
if (process.env.XML_OUTPUT_FILE) {
  const xmlDirectory = dirname(process.env.XML_OUTPUT_FILE);
  preopens[xmlDirectory] = xmlDirectory;
}

// WASI provides the C runtime environment: stdio, filesystem, clocks.
const wasi = new WASI({
  version: 'preview1',
  args: [__IREE_WASM_BINARY, ...process.argv.slice(2)],
  env: process.env,
  preopens,
});

// Start with WASI imports and merge companion imports on top.
// createWasmImports is generated by the bundler — it returns an object keyed
// by wasm import module name. For tests with no companions, it returns {}.
const imports = wasi.getImportObject();
const context = {
  memory: null
};
const companionImports = createWasmImports(context);
for (const [moduleName, moduleImports] of Object.entries(companionImports)) {
  if (imports[moduleName]) {
    Object.assign(imports[moduleName], moduleImports);
  } else {
    imports[moduleName] = moduleImports;
  }
}

// Compile and instantiate.
const wasmBytes = readFileSync(wasmPath);
const {instance} = await WebAssembly.instantiate(wasmBytes, imports);

// Wire up exported memory for companions that need it (e.g., reading
// wasm linear memory for ring drain operations).
context.memory = instance.exports.memory;

// WASI.start() calls _start() on the wasm instance. The WASI runtime handles
// proc_exit internally and sets process.exitCode.
try {
  wasi.start(instance);
} catch (error) {
  // WASI proc_exit throws to unwind the wasm call stack. process.exitCode is
  // already set by the WASI runtime. If it was not set, this is a genuine
  // error (not a proc_exit).
  if (!process.exitCode) {
    console.error(error);
    process.exitCode = 1;
  }
}