Transform JSX to Lazy Requires instead of Wrappers (#30433)

This ensures that we can keep overriding what runtime to use by
resetting modules while still using the automatic JSX plugin. This is
like the "inline requires" transform but just for JSX.

I got sick of trying to figure out workarounds to hide the extra stack
frame that appears due to the wrappers.

scripts/babel/__tests__/transform-lazy-jsx-import-test.js

@@ -0,0 +1,30 @@
+/**
+ * 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.
+ */
+'use strict';
+
+describe('transform-lazy-jsx-import', () => {
+  it('should use the mocked version of the "react" runtime in jsx', () => {
+    jest.resetModules();
+    const mock = jest.fn(type => 'fakejsx: ' + type);
+    if (__DEV__) {
+      jest.mock('react/jsx-dev-runtime', () => {
+        return {
+          jsxDEV: mock,
+        };
+      });
+    } else {
+      jest.mock('react/jsx-runtime', () => ({
+        jsx: mock,
+        jsxs: mock,
+      }));
+    }
+    // eslint-disable-next-line react/react-in-jsx-scope
+    const x = <div />;
+    expect(x).toBe('fakejsx: div');
+    expect(mock).toHaveBeenCalledTimes(1);
+  });
+});

scripts/babel/transform-lazy-jsx-import.js

@@ -0,0 +1,92 @@
+/**
+ * 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.
+ */
+'use strict';
+
+// Most of our tests call jest.resetModules in a beforeEach and the
+// re-require all the React modules. However, the JSX runtime is injected by
+// the compiler, so those bindings don't get updated. This causes warnings
+// logged by the JSX runtime to not have a component stack, because component
+// stack relies on the the secret internals object that lives on the React
+// module, which because of the resetModules call is longer the same one.
+//
+// To workaround this issue, use a transform that calls require() again before
+// every JSX invocation.
+//
+// Longer term we should migrate all our tests away from using require() and
+// resetModules, and use import syntax instead so this kind of thing doesn't
+// happen.
+
+module.exports = function replaceJSXImportWithLazy(babel) {
+  const {types: t} = babel;
+
+  function getInlineRequire(moduleName) {
+    return t.callExpression(t.identifier('require'), [
+      t.stringLiteral(moduleName),
+    ]);
+  }
+
+  return {
+    visitor: {
+      CallExpression: function (path, pass) {
+        let callee = path.node.callee;
+        if (callee.type === 'SequenceExpression') {
+          callee = callee.expressions[callee.expressions.length - 1];
+        }
+        if (callee.type === 'Identifier') {
+          // Sometimes we seem to hit this before the imports are transformed
+          // into requires and so we hit this case.
+          switch (callee.name) {
+            case '_jsxDEV':
+              path.node.callee = t.memberExpression(
+                getInlineRequire('react/jsx-dev-runtime'),
+                t.identifier('jsxDEV')
+              );
+              return;
+            case '_jsx':
+              path.node.callee = t.memberExpression(
+                getInlineRequire('react/jsx-runtime'),
+                t.identifier('jsx')
+              );
+              return;
+            case '_jsxs':
+              path.node.callee = t.memberExpression(
+                getInlineRequire('react/jsx-runtime'),
+                t.identifier('jsxs')
+              );
+              return;
+          }
+          return;
+        }
+        if (callee.type !== 'MemberExpression') {
+          return;
+        }
+        if (callee.property.type !== 'Identifier') {
+          // Needs to be jsx, jsxs, jsxDEV.
+          return;
+        }
+        if (callee.object.type !== 'Identifier') {
+          // Needs to be _reactJsxDevRuntime or _reactJsxRuntime.
+          return;
+        }
+        // Replace the cached identifier with a new require call.
+        // Relying on the identifier name is a little flaky. Should ideally pick
+        // this from the import. For some reason it sometimes has the react prefix
+        // and other times it doesn't.
+        switch (callee.object.name) {
+          case '_reactJsxDevRuntime':
+          case '_jsxDevRuntime':
+            callee.object = getInlineRequire('react/jsx-dev-runtime');
+            return;
+          case '_reactJsxRuntime':
+          case '_jsxRuntime':
+            callee.object = getInlineRequire('react/jsx-runtime');
+            return;
+        }
+      },
+    },
+  };
+};

scripts/jest/devtools/setupEnv.js

@@ -40,48 +40,3 @@ global._test_react_version_focus = (range, testName, callback) => {
 global._test_ignore_for_react_version = (testName, callback) => {
   test.skip(testName, callback);
 };
-
-// Most of our tests call jest.resetModules in a beforeEach and the
-// re-require all the React modules. However, the JSX runtime is injected by
-// the compiler, so those bindings don't get updated. This causes warnings
-// logged by the JSX runtime to not have a component stack, because component
-// stack relies on the the secret internals object that lives on the React
-// module, which because of the resetModules call is longer the same one.
-//
-// To workaround this issue, we use a proxy that re-requires the latest
-// JSX Runtime from the require cache on every function invocation.
-//
-// Longer term we should migrate all our tests away from using require() and
-// resetModules, and use import syntax instead so this kind of thing doesn't
-// happen.
-if (semver.gte(ReactVersionTestingAgainst, '17.0.0')) {
-  lazyRequireFunctionExports('react/jsx-dev-runtime');
-
-  // TODO: We shouldn't need to do this in the production runtime, but until
-  // we remove string refs they also depend on the shared state object. Remove
-  // once we remove string refs.
-  lazyRequireFunctionExports('react/jsx-runtime');
-}
-
-function lazyRequireFunctionExports(moduleName) {
-  jest.mock(moduleName, () => {
-    return new Proxy(jest.requireActual(moduleName), {
-      get(originalModule, prop) {
-        // If this export is a function, return a wrapper function that lazily
-        // requires the implementation from the current module cache.
-        if (typeof originalModule[prop] === 'function') {
-          // eslint-disable-next-line no-eval
-          const wrapper = eval(`
-            (function () {
-              return jest.requireActual(moduleName)[prop].apply(this, arguments);
-            })
-            // We use this to trick the filtering of Flight to exclude this frame.
-            //# sourceURL=<anonymous>`);
-          return wrapper;
-        } else {
-          return originalModule[prop];
-        }
-      },
-    });
-  });
-}

scripts/jest/preprocessor.js

@@ -25,6 +25,9 @@ const pathToTransformTestGatePragma = require.resolve(
 const pathToTransformReactVersionPragma = require.resolve(
   '../babel/transform-react-version-pragma'
 );
+const pathToTransformLazyJSXImport = require.resolve(
+  '../babel/transform-lazy-jsx-import'
+);
 const pathToBabelrc = path.join(__dirname, '..', '..', 'babel.config.js');
 const pathToErrorCodes = require.resolve('../error-codes/codes.json');
 
@@ -93,6 +96,8 @@ module.exports = {
         );
       }
 
+      plugins.push(pathToTransformLazyJSXImport);
+
       let sourceAst = hermesParser.parse(src, {babel: true});
       return {
         code: babel.transformFromAstSync(
@@ -122,6 +127,7 @@ module.exports = {
       pathToTransformInfiniteLoops,
       pathToTransformTestGatePragma,
       pathToTransformReactVersionPragma,
+      pathToTransformLazyJSXImport,
       pathToErrorCodes,
     ],
     [

scripts/jest/setupTests.js

@@ -264,46 +264,3 @@ if (process.env.REACT_CLASS_EQUIVALENCE_TEST) {
     return fn(flags);
   };
 }
-
-// Most of our tests call jest.resetModules in a beforeEach and the
-// re-require all the React modules. However, the JSX runtime is injected by
-// the compiler, so those bindings don't get updated. This causes warnings
-// logged by the JSX runtime to not have a component stack, because component
-// stack relies on the the secret internals object that lives on the React
-// module, which because of the resetModules call is longer the same one.
-//
-// To workaround this issue, we use a proxy that re-requires the latest
-// JSX Runtime from the require cache on every function invocation.
-//
-// Longer term we should migrate all our tests away from using require() and
-// resetModules, and use import syntax instead so this kind of thing doesn't
-// happen.
-lazyRequireFunctionExports('react/jsx-dev-runtime');
-
-// TODO: We shouldn't need to do this in the production runtime, but until
-// we remove string refs they also depend on the shared state object. Remove
-// once we remove string refs.
-lazyRequireFunctionExports('react/jsx-runtime');
-
-function lazyRequireFunctionExports(moduleName) {
-  jest.mock(moduleName, () => {
-    return new Proxy(jest.requireActual(moduleName), {
-      get(originalModule, prop) {
-        // If this export is a function, return a wrapper function that lazily
-        // requires the implementation from the current module cache.
-        if (typeof originalModule[prop] === 'function') {
-          // eslint-disable-next-line no-eval
-          const wrapper = eval(`
-            (function () {
-              return jest.requireActual(moduleName)[prop].apply(this, arguments);
-            })
-            // We use this to trick the filtering of Flight to exclude this frame.
-            //# sourceURL=<anonymous>`);
-          return wrapper;
-        } else {
-          return originalModule[prop];
-        }
-      },
-    });
-  });
-}