Stop cloning the AST in the transform worker's main Babel pass (#2034)

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

@@ -24,9 +24,43 @@ jest
   .mock('metro-transform-plugins', () => ({
     ...jest.requireActual('metro-transform-plugins'),
     inlinePlugin: () => ({}),
-    constantFoldingPlugin: () => ({}),
   }))
-  .mock('metro-minify-terser');
+  .mock('metro-minify-terser')
+  .mock(
+    'stale-scope-babel-transformer',
+    () => {
+      const {transformSync, types} = jest.requireActual('@babel/core');
+      return {
+        // Replaces the second statement with a call to `stringifySafe` without
+        // updating the scope, which Babel crawled before the replacement.
+        transform({src}: {src: string}) {
+          return transformSync(src, {
+            ast: true,
+            babelrc: false,
+            code: false,
+            configFile: false,
+            plugins: [
+              () => ({
+                visitor: {
+                  Program: {
+                    exit(path) {
+                      path.node.body[1] = types.expressionStatement(
+                        types.callExpression(
+                          types.identifier('stringifySafe'),
+                          [types.stringLiteral('value')],
+                        ),
+                      );
+                    },
+                  },
+                },
+              }),
+            ],
+          });
+        },
+      };
+    },
+    {virtual: true},
+  );
 
 import type {JsTransformerConfig, JsTransformOptions} from '../index';
 import typeof * as TransformerType from '../index';
@@ -43,6 +77,10 @@ const babelTransformerPath =
 const HEADER_DEV =
   '__d(function (global, require, _$$_IMPORT_DEFAULT, _$$_IMPORT_ALL, module, exports, _dependencyMap) {';
 const HEADER_PROD = '__d(function (g, r, i, a, m, e, d) {';
+const STALE_SCOPE_SOURCE = [
+  'function stringifySafe(arg) { return String(arg); }',
+  'STALE_SCOPE_REFERENCE;',
+].join('\n');
 
 let fs: FSType;
 let Transformer: TransformerType;
@@ -203,6 +241,43 @@ test('transforms a module with dependencies', async () => {
   ]);
 });
 
+test('keeps functions referenced by earlier passes in production', async () => {
+  const result = await Transformer.transform(
+    {...baseConfig, babelTransformerPath: 'stale-scope-babel-transformer'},
+    '/root',
+    'local/file.js',
+    Buffer.from(STALE_SCOPE_SOURCE, 'utf8'),
+    {...baseTransformOptions, dev: false, minify: true},
+  );
+
+  expect(result.output[0].data.code).toContain('function stringifySafe(arg)');
+  expect(result.output[0].data.code).toContain('stringifySafe("value")');
+});
+
+test('Babel keeps stale scope bindings when reusing an AST', () => {
+  const {ast} = jest
+    .requireMock('stale-scope-babel-transformer')
+    .transform({src: STALE_SCOPE_SOURCE});
+  const {transformFromAstSync} = jest.requireActual('@babel/core');
+  const {constantFoldingPlugin} = jest.requireActual('metro-transform-plugins');
+
+  const result = transformFromAstSync(ast, '', {
+    ast: false,
+    babelrc: false,
+    cloneInputAst: false,
+    code: true,
+    configFile: false,
+    plugins: [constantFoldingPlugin],
+  });
+
+  // Babel reuses the scope crawled before `stringifySafe` was called, so
+  // constant folding deletes the function while its call remains. This is why
+  // transformJS clears the scope cache before constant folding. If this test
+  // starts failing, Babel no longer reuses stale scopes: delete that
+  // `traverse.cache.clearScope()` call and this test.
+  expect(result.code).toBe('stringifySafe("value");');
+});
+
 test('transforms an es module with asyncToGenerator', async () => {
   const result = await Transformer.transform(
     baseConfig,

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';
@@ -342,12 +342,7 @@ async function transformJS(
     transformFromAstSync(ast, '', {
       ast: true,
       babelrc: false,
-      // Not-Cloning the input AST here should be safe because other code paths above this call
-      // are mutating the AST as well and no code is depending on the original AST.
-      // 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: false,
       code: false,
       comments: true,
       configFile: false,
@@ -361,6 +356,11 @@ 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).
+    // 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). Remove this clear when the
+    // "Babel keeps stale scope bindings when reusing an AST" test fails.
+    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 |',
+  ]);
+});