@bytecodealliance/jco/src/cmd/opt.js

JavaScript tooling for working with WebAssembly Components

import { $init, tools } from "../../obj/wasm-tools.js";
const { metadataShow, print } = tools;
import { writeFile } from "fs/promises";
import { fileURLToPath } from "url";
import c from "chalk-template";
import {
  readFile,
  sizeStr,
  fixedDigitDisplay,
  table,
  spawnIOTmp,
  setShowSpinner,
  getShowSpinner,
} from "../common.js";
import ora from "#ora";

export async function opt(componentPath, opts, program) {
  await $init;
  const varIdx = program.parent.rawArgs.indexOf("--");
  if (varIdx !== -1) opts.optArgs = program.parent.rawArgs.slice(varIdx + 1);
  const componentBytes = await readFile(componentPath);

  if (!opts.quiet) setShowSpinner(true);
  const optPromise = optimizeComponent(componentBytes, opts);
  const { component, compressionInfo } = await optPromise;

  await writeFile(opts.output, component);

  let totalBeforeBytes = 0,
    totalAfterBytes = 0;

  if (!opts.quiet)
    console.log(c`
{bold Optimized WebAssembly Component Internal Core Modules:}

${table(
  [
    ...compressionInfo.map(({ beforeBytes, afterBytes }, i) => {
      totalBeforeBytes += beforeBytes;
      totalAfterBytes += afterBytes;
      return [
        ` - Core Module ${i + 1}:  `,
        sizeStr(beforeBytes),
        " -> ",
        c`{cyan ${sizeStr(afterBytes)}} `,
        `(${fixedDigitDisplay((afterBytes / beforeBytes) * 100, 2)}%)`,
      ];
    }),
    ["", "", "", "", ""],
    [
      ` = Total:  `,
      `${sizeStr(totalBeforeBytes)}`,
      ` => `,
      c`{cyan ${sizeStr(totalAfterBytes)}} `,
      `(${fixedDigitDisplay((totalAfterBytes / totalBeforeBytes) * 100, 2)}%)`,
    ],
  ],
  [, , , , "right"]
)}`);
}

/**
 * Counts the byte length for the LEB128 encoding of a number.
 * @param {number} val
 * @returns {number}
 */
function byteLengthLEB128(val) {
  let len = 0;
  do {
    val >>>= 7;
    len++;
  } while (val !== 0);
  return len;
}

/**
 *
 * @param {Uint8Array} componentBytes
 * @param {{ quiet: boolean, asyncify?: boolean, optArgs?: string[], noVerify?: boolean }} opts?
 * @returns {Promise<{ component: Uint8Array, compressionInfo: { beforeBytes: number, afterBytes: number }[] >}
 */
export async function optimizeComponent(componentBytes, opts) {
  await $init;
  const showSpinner = getShowSpinner();
  let spinner;
  try {
    let componentMetadata = metadataShow(componentBytes);
    componentMetadata.forEach((metadata, index) => {
        // add index to the metadata object
        metadata.index = index;
        const size = metadata.range[1] - metadata.range[0];
        // compute previous LEB128 encoding length
        metadata.prevLEBLen = byteLengthLEB128(size);
    });
    const coreModules = componentMetadata
      .filter(({ metaType }) => metaType.tag === "module");

    // log number of core Wasm modules to be run with wasm-opt
    let completed = 0;
    const spinnerText = () =>
      c`{cyan ${completed} / ${coreModules.length}} Running Binaryen on WebAssembly Component Internal Core Modules \n`;
    if (showSpinner) {
      spinner = ora({
        color: "cyan",
        spinner: "bouncingBar",
      }).start();
      spinner.text = spinnerText();
    }

    // gather the options for wasm-opt. optionally, adding the asyncify flag
    const args = opts?.optArgs ?
      [...opts.optArgs] :
      ['-Oz', '--low-memory-unused', '--enable-bulk-memory', '--strip-debug'];
    if (opts?.asyncify) args.push('--asyncify');


    // process core Wasm modules with wasm-opt
    await Promise.all(
      coreModules.map(async (metadata) => {
        if (metadata.metaType.tag === "module") {
          // store the wasm-opt processed module in the metadata 
          metadata.optimized = await wasmOpt(
            componentBytes.subarray(metadata.range[0], metadata.range[1]),
            args);

          // compute the size change, including the change to
          // the LEB128 encoding of the size change
          const prevModuleSize = metadata.range[1] - metadata.range[0];
          const newModuleSize = metadata.optimized.byteLength;
          metadata.newLEBLen = byteLengthLEB128(newModuleSize);
          metadata.sizeChange = newModuleSize - prevModuleSize;

          if (spinner) {
            completed++;
            spinner.text = spinnerText();
          }
        }
      }));

    // organize components in modules into tree parent and children
    const nodes = componentMetadata.slice(1);
    const getChildren = (parentIndex) => {
      const children = [];
      for (let i = 0; i < nodes.length; i++) {
        const metadata = nodes[i];
        if (metadata.parentIndex === parentIndex) {
          nodes.splice(i, 1); // remove from nodes
          i--;
          metadata.children = getChildren(metadata.index);
          metadata.sizeChange = metadata.children
            .reduce((total, {prevLEBLen, newLEBLen, sizeChange}) => {
              return sizeChange ?
                total + sizeChange + newLEBLen - prevLEBLen :
                total;
            },
            metadata.sizeChange || 0);
          const prevSize = metadata.range[1] - metadata.range[0];
          metadata.newLEBLen = byteLengthLEB128(prevSize + metadata.sizeChange);
          children.push(metadata);
        }
      }
      return children;
    };
    const componentTree = getChildren(0);

    // compute the total size change in the component binary
    const sizeChange = componentTree
      .reduce((total, {prevLEBLen, newLEBLen, sizeChange}) => {
        return total + (sizeChange || 0) + newLEBLen - prevLEBLen;
      }, 0);

    let outComponentBytes = new Uint8Array(
      componentBytes.byteLength + sizeChange);
    let nextReadPos = 0, nextWritePos = 0;

    const write = ({prevLEBLen, range, optimized, children, sizeChange}) => {
      // write from the last read to the LEB byte start
      outComponentBytes.set(
        componentBytes.subarray(nextReadPos, range[0] - prevLEBLen),
        nextWritePos
      );
      nextWritePos += range[0] - prevLEBLen - nextReadPos;

      // write the new LEB bytes
      let val = range[1] - range[0] + sizeChange;
      do {
        const byte = val & 0x7f;
        val >>>= 7;
        outComponentBytes[nextWritePos++] = val === 0 ? byte : byte | 0x80;
      } while (val !== 0);

      if (optimized) {
        // write the core module
        outComponentBytes.set(optimized, nextWritePos);
        nextReadPos = range[1];
        nextWritePos += optimized.byteLength;
      } else if (children.length > 0) {
        // write child components / modules
        nextReadPos = range[0];
        children.forEach(write);
      } else {
        // write component
        outComponentBytes.set(
          componentBytes.subarray(range[0], range[1]),
          nextWritePos
        );
        nextReadPos = range[1];
        nextWritePos += range[1] - range[0];
      }
    };

    // write each top-level component / module
    componentTree.forEach(write);

    // write remaining
    outComponentBytes.set(
      componentBytes.subarray(nextReadPos),
      nextWritePos
    );

    // verify it still parses ok
    if (!opts?.noVerify) {
      try {
        await print(outComponentBytes);
      } catch (e) {
        throw new Error(
          `Internal error performing optimization.\n${e.message}`
        );
      }
    }

    return {
      component: outComponentBytes,
      compressionInfo: coreModules.map(({range, optimized}) => ({
        beforeBytes: range[1] - range[0],
        afterBytes: optimized?.byteLength,
      })),
    };
  } finally {
    if (spinner) spinner.stop();
  }
}

/**
 * @param {Uint8Array} source
 * @param {Array<string>} args
 * @returns {Promise<Uint8Array>}
 */
async function wasmOpt(source, args) {
  const wasmOptPath = fileURLToPath(
    import.meta.resolve("binaryen/bin/wasm-opt")
  );

  try {
    return await spawnIOTmp(wasmOptPath, source, [...args, "-o"]);
  } catch (e) {
    if (e.toString().includes("BasicBlock requested"))
      return wasmOpt(source, args);
    throw e;
  }
}