Transform JSX to Lazy Requires instead of 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];
-        }
-      },
-    });
-  });
-}