Experiment: skip AST cloning in the transform worker (#2033)

packages/metro-config/src/defaults/index.js

@@ -136,6 +136,7 @@ const getDefaultValues = (projectRoot: ?string): ConfigT => ({
     unstable_disableModuleWrapping: false,
     unstable_disableNormalizePseudoGlobals: false,
     unstable_compactOutput: false,
+    unstable_disableInputAstCloning: false,
     unstable_memoizeInlineRequires: false,
     unstable_workerThreads: false,
   },

packages/metro-transform-worker/API.md

@@ -41,6 +41,7 @@ export type JsTransformerConfig = Readonly<{
   unstable_disableModuleWrapping: boolean;
   unstable_disableNormalizePseudoGlobals: boolean;
   unstable_compactOutput: boolean;
+  unstable_disableInputAstCloning?: boolean | undefined;
   unstable_allowRequireContext: boolean;
   unstable_memoizeInlineRequires?: boolean | undefined;
   unstable_nonMemoizedInlineRequires?: ReadonlyArray<string> | undefined;

packages/metro-transform-worker/src/index.js

@@ -37,7 +37,7 @@ import type {
 
 import * as assetTransformer from './utils/assetTransformer';
 import getMinifier from './utils/getMinifier';
-import {transformFromAstSync} from '@babel/core';
+import {transformFromAstSync, traverse} from '@babel/core';
 import generate from '@babel/generator';
 import * as babylon from '@babel/parser';
 import * as types from '@babel/types';
@@ -107,6 +107,8 @@ export type JsTransformerConfig = Readonly<{
   unstable_disableModuleWrapping: boolean,
   unstable_disableNormalizePseudoGlobals: boolean,
   unstable_compactOutput: boolean,
+  /** Skip cloning the AST before the main Babel pass. */
+  unstable_disableInputAstCloning?: boolean,
   /** Enable `require.context` statements which can be used to import multiple files in a directory. */
   unstable_allowRequireContext: boolean,
   /** With inlineRequires, enable a module-scope memo var and inline as (v || v=require('foo')) */
@@ -347,7 +349,7 @@ async function transformJS(
       // However, switching the flag to false caused issues with ES Modules if `experimentalImportSupport` isn't used https://github.com/react/metro/issues/641
       // either because one of the plugins is doing something funky or Babel messes up some caches.
       // Make sure to test the above mentioned case before flipping the flag back to false.
-      cloneInputAst: true,
+      cloneInputAst: config.unstable_disableInputAstCloning !== true,
       code: false,
       comments: true,
       configFile: false,
@@ -361,6 +363,12 @@ async function transformJS(
     // Run the constant folding plugin in its own pass, avoiding race conditions
     // with other plugins that have exit() visitors on Program (e.g. the ESM
     // transform).
+    if (config.unstable_disableInputAstCloning === true) {
+      // Babel reuses scopes cached by earlier passes over the same uncloned
+      // AST, so constant folding would strip functions those passes started
+      // using (https://github.com/react/metro/issues/641).
+      traverse.cache.clearScope();
+    }
     ast = nullthrows(
       transformFromAstSync(ast, '', {
         ast: true,

scripts/benchmark-patches/.gitignore

@@ -0,0 +1 @@
+dist/

scripts/benchmark-patches/__tests__/benchmark-patches-test.js

@@ -0,0 +1,57 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @format
+ * @oncall react_native
+ */
+
+'use strict';
+
+import {rotate} from '../benchmark-patches';
+import {spawnSync} from 'node:child_process';
+import path from 'node:path';
+
+const SCRIPT = path.join(__dirname, '..', 'benchmark-patches.js');
+
+function runScript(...args) {
+  return spawnSync(process.execPath, [SCRIPT, ...args], {encoding: 'utf8'});
+}
+
+test('prints help without arguments and with --help', () => {
+  for (const args of [[], ['--help']]) {
+    const {status, stdout} = runScript(...args);
+    expect(status).toBe(0);
+    expect(stdout).toMatch(/^Usage: node benchmark-patches\.js <benchmark>/);
+  }
+});
+
+test('rejects invalid arguments', () => {
+  const noBenchmark = runScript('Entry.js:ios');
+  expect(noBenchmark.status).not.toBe(0);
+  expect(noBenchmark.stderr).toContain(
+    'Unknown benchmark "Entry.js:ios", expected one of: bundling_speed',
+  );
+
+  const badRuns = runScript('bundling_speed', '--runs=1', 'Entry.js:ios');
+  expect(badRuns.status).not.toBe(0);
+  expect(badRuns.stderr).toContain('--runs must be an integer >= 2, got "1"');
+
+  const noTargets = runScript('bundling_speed');
+  expect(noTargets.status).not.toBe(0);
+  expect(noTargets.stderr).toContain('Missing targets for bundling_speed');
+
+  const noPlatform = runScript('bundling_speed', 'Entry.js');
+  expect(noPlatform.status).not.toBe(0);
+  expect(noPlatform.stderr).toContain('Missing platform in "Entry.js"');
+});
+
+test('rotates every variant through every position', () => {
+  expect([0, 1, 2].map(offset => rotate(['a', 'b', 'c'], offset))).toEqual([
+    ['a', 'b', 'c'],
+    ['b', 'c', 'a'],
+    ['c', 'a', 'b'],
+  ]);
+});

scripts/benchmark-patches/__tests__/bundling_speed-test.js

@@ -0,0 +1,48 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @format
+ * @oncall react_native
+ */
+
+'use strict';
+
+import {formatTable, parseTargets} from '../benchmarks/bundling_speed';
+
+test('parses <entry>:<platform> into a prod and a dev bundle', () => {
+  expect(parseTargets(['src/App.js:ios'], {}).map(b => b.name)).toEqual([
+    'App-ios-prod',
+    'App-ios-dev',
+  ]);
+  expect(() => parseTargets(['Entry.js'], {})).toThrow(
+    'Missing platform in "Entry.js"',
+  );
+});
+
+test('formats means, deltas with t, and byte identity of every run', () => {
+  const samples = (walls, cpus, hashes) =>
+    walls.map((wall, i) => ({cpu: cpus[i], hash: hashes[i], size: 1000, wall}));
+  const results = new Map([
+    [
+      'App-ios-prod',
+      new Map([
+        ['unpatched', samples([10, 12], [100, 120], ['a', 'b'])],
+        ['faster', samples([8, 10], [80, 100], ['a', 'b'])],
+        ['broken', samples([10, 12], [100, 120], ['a', 'x'])],
+      ]),
+    ],
+  ]);
+  const table = formatTable(
+    [{name: 'App-ios-prod'}],
+    [{name: 'unpatched'}, {name: 'faster'}, {name: 'broken'}],
+    results,
+  );
+  expect(table.split('\n').slice(2, 5)).toEqual([
+    '| App-ios-prod | unpatched | 11.0 ± 1.4 |  | 110.0 ± 14.1 |  | 1,000 | yes |',
+    '| App-ios-prod | faster | 9.0 ± 1.4 | -18.18% (t -1.4) | 90.0 ± 14.1 | -18.18% (t -1.4) | 1,000 | yes |',
+    '| App-ios-prod | broken | 11.0 ± 1.4 | +0.00% (t 0.0) | 110.0 ± 14.1 | +0.00% (t 0.0) | 1,000 | NO |',
+  ]);
+});

scripts/benchmark-patches/benchmark-patches.js

@@ -0,0 +1,182 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @format
+ * @oncall react_native
+ */
+
+'use strict';
+
+const bundlingSpeed = require('./benchmarks/bundling_speed');
+const run = require('./run');
+const fs = require('node:fs');
+const path = require('node:path');
+const {parseArgs} = require('node:util');
+
+const BENCHMARKS = {bundling_speed: bundlingSpeed};
+
+const DIST_DIR = path.join(__dirname, 'dist');
+const PATCHES_DIR = path.join(__dirname, 'patches');
+const METRO_ROOT = path.resolve(__dirname, '../..');
+
+const HELP = `Usage: node benchmark-patches.js <benchmark> [--runs=N] <target>...
+
+Runs <benchmark> on each target, unpatched and with every *.patch in the
+patches directory applied, and prints a table comparing each patch with
+unpatched.
+
+Patches directory:
+  ${PATCHES_DIR}
+
+Each target gets one unmeasured warmup run, then variants are interleaved in a
+rotating order. dist/ is cleared on start and holds the last run's outputs
+(dist/<variant>/) and the table (dist/results.md).
+
+Options:
+  --runs=N  Measured runs per variant, at least 2 (default: 5).
+  --help    Show this help.
+
+Patches are diffs relative to the Metro repository root, e.g. \`git diff\`
+output.
+
+Benchmarks:
+
+${Object.values(BENCHMARKS)
+  .map(benchmark => benchmark.HELP)
+  .join('\n')}`;
+
+function patch(patchFile, {reverse = false, dryRun = false} = {}) {
+  run(
+    'patch',
+    [
+      '-p1',
+      '--batch',
+      '--silent',
+      '--no-backup-if-mismatch',
+      reverse ? '--reverse' : '--forward',
+      ...(dryRun ? ['--dry-run'] : []),
+      '-d',
+      METRO_ROOT,
+      '-i',
+      patchFile,
+    ],
+    {stdio: 'inherit'},
+  );
+}
+
+// Rotating the variant order each run spreads position bias evenly.
+function rotate(items, offset) {
+  return items.map((_, i) => items[(offset + i) % items.length]);
+}
+
+function main() {
+  const benchmarkOptions = Object.assign(
+    {},
+    ...Object.values(BENCHMARKS).map(benchmark => benchmark.OPTIONS),
+  );
+  const {values, positionals} = parseArgs({
+    allowPositionals: true,
+    options: {
+      help: {type: 'boolean'},
+      runs: {type: 'string', default: '5'},
+      ...benchmarkOptions,
+    },
+  });
+  if (values.help === true || positionals.length === 0) {
+    process.stdout.write(HELP);
+    return;
+  }
+  const [benchmarkName, ...specs] = positionals;
+  if (!Object.hasOwn(BENCHMARKS, benchmarkName)) {
+    throw new Error(
+      `Unknown benchmark "${benchmarkName}", expected one of: ` +
+        Object.keys(BENCHMARKS).join(', '),
+    );
+  }
+  const benchmark = BENCHMARKS[benchmarkName];
+  const runs = Number(values.runs);
+  if (!Number.isInteger(runs) || runs < 2) {
+    throw new Error(`--runs must be an integer >= 2, got "${values.runs}"`);
+  }
+  if (specs.length === 0) {
+    throw new Error(`Missing targets for ${benchmarkName}`);
+  }
+  const targets = benchmark.parseTargets(specs, values);
+  const variants = [
+    {name: 'unpatched', patchFile: null},
+    ...fs
+      .readdirSync(PATCHES_DIR)
+      .filter(file => file.endsWith('.patch'))
+      .sort()
+      .map(file => ({
+        name: path.basename(file, '.patch'),
+        patchFile: path.join(PATCHES_DIR, file),
+      })),
+  ];
+  for (const {patchFile} of variants) {
+    if (patchFile != null) {
+      patch(patchFile, {dryRun: true});
+    }
+  }
+
+  fs.rmSync(DIST_DIR, {force: true, recursive: true});
+  for (const variant of variants) {
+    fs.mkdirSync(path.join(DIST_DIR, variant.name), {recursive: true});
+  }
+
+  // Ctrl-C reaches the benchmark child; the finally block below reverts the
+  // patch.
+  process.on('SIGINT', () => {});
+
+  const results = new Map(
+    targets.map(target => [
+      target.name,
+      new Map(variants.map(variant => [variant.name, []])),
+    ]),
+  );
+  for (const target of targets) {
+    const warmup = benchmark.measure(
+      target,
+      path.join(DIST_DIR, variants[0].name),
+    );
+    console.error(`[warmup] ${target.name}: ${benchmark.describe(warmup)}`);
+  }
+
+  const total = runs * targets.length * variants.length;
+  let count = 0;
+  for (let i = 1; i <= runs; i++) {
+    for (const target of targets) {
+      for (const variant of rotate(variants, i - 1)) {
+        if (variant.patchFile != null) {
+          patch(variant.patchFile);
+        }
+        let result;
+        try {
+          result = benchmark.measure(target, path.join(DIST_DIR, variant.name));
+        } finally {
+          if (variant.patchFile != null) {
+            patch(variant.patchFile, {reverse: true});
+          }
+        }
+        results.get(target.name).get(variant.name).push(result);
+        console.error(
+          `[${++count}/${total}] run ${i} ${target.name} ${variant.name}: ` +
+            benchmark.describe(result),
+        );
+      }
+    }
+  }
+
+  const table = benchmark.formatTable(targets, variants, results);
+  fs.writeFileSync(path.join(DIST_DIR, 'results.md'), table);
+  process.stdout.write(table);
+}
+
+if (require.main === module) {
+  main();
+}
+
+module.exports = {BENCHMARKS, rotate};