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

/**
 * 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.
 */

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,
} from 'babel-plugin-react-compiler';
import {basename} from 'path';

/**
 * -- IMPORTANT --
 * When making changes to any babel plugins defined this file
 * (e.g. `ReactForgetFunctionTransform`), make sure to bump e2eTransformerCacheKey
 * as our script files are currently not used for babel cache breaking!!
 */
const e2eTransformerCacheKey = 1;
const forgetOptions: EnvironmentConfig = validateEnvironmentConfig({
  enableAssumeHooksFollowRulesOfReact: true,
  enableFunctionOutlining: false,
});
const debugMode = process.env['DEBUG_FORGET_COMPILER'] != null;

module.exports = (useForget: boolean) => {
  function createTransformer() {
    return babelJest.createTransformer({
      passPerPreset: true,
      presets: [
        '@babel/preset-typescript',
        {
          plugins: [
            useForget
              ? [
                  ReactForgetFunctionTransform,
                  {
                    /*
                     * 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(),
                    transformOptionsCacheKey: forgetOptions,
                    e2eTransformerCacheKey,
                  },
                ]
              : '@babel/plugin-syntax-jsx',
          ],
        },
        '@babel/preset-react',
        {
          plugins: [
            [
              function BabelPluginRewriteRequirePath(): {visitor: Visitor} {
                return {
                  visitor: {
                    CallExpression(path: NodePath<CallExpression>): void {
                      const {callee} = path.node;
                      if (
                        callee.type === 'Identifier' &&
                        callee.name === 'require'
                      ) {
                        const arg = path.node.arguments[0];
                        if (arg.type === 'StringLiteral') {
                          /*
                           * The compiler adds requires of "React", which is expected to be a wrapper
                           * around the "react" package. For tests, we just rewrite the require.
                           */
                          if (arg.value === 'React') {
                            arg.value = 'react';
                          }
                        }
                      }
                    },
                  },
                };
              },
            ],
            '@babel/plugin-transform-modules-commonjs',
          ],
        },
      ],
      targets: {
        esmodules: true,
      },
    } as any);
    /*
     * typecast needed as DefinitelyTyped does not have updated Babel configs types yet
     * (missing passPerPreset and targets).
     */
  }

  return {
    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',
        '_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,
  };
}