[compiler][be] Move e2e tests to BabelPlugin transformer (#32706)

Clean up jest-e2e setup since
#32663 and other features need
program context (e.g. changing imports)
---
[//]: # (BEGIN SAPLING FOOTER)
Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with [ReviewStack](https://reviewstack.dev/facebook/react/pull/32706).
* #32663
* __->__ #32706

Author: mofeiZ

compiler/packages/babel-plugin-react-compiler/scripts/jest/e2e-forget.config.js

@@ -7,7 +7,4 @@
 
 const makeE2EConfig = require('../jest/makeE2EConfig');
 
-const config = makeE2EConfig('e2e with forget', true);
-config.setupFilesAfterEnv = ['<rootDir>/../scripts/jest/setupEnvE2E.js'];
-
-module.exports = config;
+module.exports = makeE2EConfig('e2e with forget', true);

compiler/packages/babel-plugin-react-compiler/scripts/jest/makeTransform.ts

@@ -5,19 +5,16 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import {jsx} from '@babel/plugin-syntax-jsx';
 import babelJest from 'babel-jest';
-import {compile} from 'babel-plugin-react-compiler';
-import {execSync} from 'child_process';
-
-import type {NodePath, Visitor} from '@babel/traverse';
-import type {CallExpression, FunctionDeclaration} from '@babel/types';
-import * as t from '@babel/types';
 import {
-  EnvironmentConfig,
   validateEnvironmentConfig,
+  EnvironmentConfig,
 } from 'babel-plugin-react-compiler';
-import {basename} from 'path';
+import {execSync} from 'child_process';
+
+import type {NodePath, Visitor} from '@babel/traverse';
+import type {CallExpression} from '@babel/types';
+import BabelPluginReactCompiler from 'babel-plugin-react-compiler';
 
 /**
  * -- IMPORTANT --
@@ -28,10 +25,19 @@ import {basename} from 'path';
 const e2eTransformerCacheKey = 1;
 const forgetOptions: EnvironmentConfig = validateEnvironmentConfig({
   enableAssumeHooksFollowRulesOfReact: true,
-  enableFunctionOutlining: false,
 });
 const debugMode = process.env['DEBUG_FORGET_COMPILER'] != null;
 
+const compilerCacheKey = execSync(
+  'yarn --silent --cwd ../.. hash packages/babel-plugin-react-compiler/dist',
+)
+  .toString()
+  .trim();
+
+if (debugMode) {
+  console.log('cachebreaker', compilerCacheKey);
+}
+
 module.exports = (useForget: boolean) => {
   function createTransformer() {
     return babelJest.createTransformer({
@@ -42,15 +48,14 @@ module.exports = (useForget: boolean) => {
           plugins: [
             useForget
               ? [
-                  ReactForgetFunctionTransform,
+                  BabelPluginReactCompiler,
                   {
+                    environment: forgetOptions,
                     /*
                      * Jest hashes the babel config as a cache breaker.
                      * (see https://github.com/jestjs/jest/blob/v29.6.2/packages/babel-jest/src/index.ts#L84)
                      */
-                    compilerCacheKey: execSync(
-                      'yarn --silent --cwd ../.. hash packages/babel-plugin-react-compiler/dist',
-                    ).toString(),
+                    compilerCacheKey,
                     transformOptionsCacheKey: forgetOptions,
                     e2eTransformerCacheKey,
                   },
@@ -105,104 +110,3 @@ module.exports = (useForget: boolean) => {
     createTransformer,
   };
 };
-
-// Mostly copied from react/scripts/babel/transform-forget.js
-function isReactComponentLike(fn: NodePath<FunctionDeclaration>): boolean {
-  let isReactComponent = false;
-  let hasNoUseForgetDirective = false;
-
-  /*
-   * React components start with an upper case letter,
-   * React hooks start with `use`
-   */
-  if (
-    fn.node.id == null ||
-    (fn.node.id.name[0].toUpperCase() !== fn.node.id.name[0] &&
-      !/^use[A-Z0-9]/.test(fn.node.id.name))
-  ) {
-    return false;
-  }
-
-  fn.traverse({
-    DirectiveLiteral(path) {
-      if (path.node.value === 'use no forget') {
-        hasNoUseForgetDirective = true;
-      }
-    },
-
-    JSX(path) {
-      // Is there is a JSX node created in the current function context?
-      if (path.scope.getFunctionParent()?.path.node === fn.node) {
-        isReactComponent = true;
-      }
-    },
-
-    CallExpression(path) {
-      // Is there hook usage?
-      if (
-        path.node.callee.type === 'Identifier' &&
-        !/^use[A-Z0-9]/.test(path.node.callee.name)
-      ) {
-        isReactComponent = true;
-      }
-    },
-  });
-
-  if (hasNoUseForgetDirective) {
-    return false;
-  }
-
-  return isReactComponent;
-}
-
-function ReactForgetFunctionTransform() {
-  const compiledFns = new Set();
-  const visitor = {
-    FunctionDeclaration(fn: NodePath<FunctionDeclaration>, state: any): void {
-      if (compiledFns.has(fn.node)) {
-        return;
-      }
-
-      if (!isReactComponentLike(fn)) {
-        return;
-      }
-      if (debugMode) {
-        const filename = basename(state.file.opts.filename);
-        if (fn.node.loc && fn.node.id) {
-          console.log(
-            ` Compiling ${filename}:${fn.node.loc.start.line}:${fn.node.loc.start.column}  ${fn.node.id.name}`,
-          );
-        } else {
-          console.log(` Compiling ${filename} ${fn.node.id?.name}`);
-        }
-      }
-
-      const compiled = compile(
-        fn,
-        forgetOptions,
-        'Other',
-        'all_features',
-        '_c',
-        null,
-        null,
-        null,
-      );
-      compiledFns.add(compiled);
-
-      const fun = t.functionDeclaration(
-        compiled.id,
-        compiled.params,
-        compiled.body,
-        compiled.generator,
-        compiled.async,
-      );
-      fn.replaceWith(fun);
-      fn.skip();
-    },
-  };
-  return {
-    name: 'react-forget-e2e',
-    inherits: jsx,
-    visitor,
-  };
-}

compiler/packages/babel-plugin-react-compiler/scripts/jest/setupEnvE2E.js

@@ -12,6 +12,7 @@ globalThis.constantValue = 'global test value';
 
 test('literal-constant-propagation', () => {
   function Component() {
+    'use memo';
     const x = 'test value 1';
     return <div>{x}</div>;
   }
@@ -38,6 +39,7 @@ test('literal-constant-propagation', () => {
 
 test('global-constant-propagation', () => {
   function Component() {
+    'use memo';
     const x = constantValue;
 
     return <div>{x}</div>;
@@ -65,6 +67,7 @@ test('global-constant-propagation', () => {
 
 test('lambda-constant-propagation', () => {
   function Component() {
+    'use memo';
     const x = 'test value 1';
     const getDiv = () => <div>{x}</div>;
     return getDiv();
@@ -92,6 +95,7 @@ test('lambda-constant-propagation', () => {
 
 test('lambda-constant-propagation-of-phi-node', () => {
   function Component({noopCallback}) {
+    'use memo';
     const x = 'test value 1';
     if (constantValue) {
       noopCallback();

compiler/packages/babel-plugin-react-compiler/src/__tests__/e2e/constant-prop.e2e.js

@@ -16,6 +16,7 @@ function Button({label}) {
 
 let currentTheme = 'light';
 function useTheme() {
+  'use memo';
   return currentTheme;
 }
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/e2e/update-button.e2e.js

@@ -10,6 +10,7 @@ import * as React from 'react';
 import {expectLogsAndClear, log} from './expectLogs';
 
 function Counter(props) {
+  'use memo';
   let value = props.value;
   let a = value++;
   expect(a).toBe(props.value); // postfix