1/**2 * Copyright (c) Meta Platforms, Inc. and affiliates.3 *4 * This source code is licensed under the MIT license found in the5 * LICENSE file in the root directory of this source tree.6 */78import type * as BabelCore from '@babel/core';9import {NodePath} from '@babel/core';10import * as t from '@babel/types';1112export default function AnnotateReactCodeBabelPlugin(13  _babel: typeof BabelCore,14): BabelCore.PluginObj {15  return {16    name: 'annotate-react-code',17    visitor: {18      Program(prog): void {19        annotate(prog);20      },21    },22  };23}2425function annotate(program: NodePath<t.Program>): void {26  function traverseFn(fn: BabelFn): void {27    if (!shouldVisit(fn)) {28      return;29    }3031    fn.skip();3233    const body = fn.node.body;34    if (t.isBlockStatement(body)) {35      body.body.unshift(buildTypeOfReactForget());36    }37  }3839  program.traverse({40    FunctionDeclaration: traverseFn,41    FunctionExpression: traverseFn,42    ArrowFunctionExpression: traverseFn,43  });44}4546function shouldVisit(fn: BabelFn): boolean {47  return (48    // Component declarations are known components49    (fn.isFunctionDeclaration() && isComponentDeclaration(fn.node)) ||50    // Otherwise check if this is a component or hook-like function51    isComponentOrHookLike(fn)52  );53}5455function buildTypeOfReactForget(): t.Statement {56  // typeof globalThis[Symbol.for("react_forget")]57  return t.expressionStatement(58    t.unaryExpression(59      'typeof',60      t.memberExpression(61        t.identifier('globalThis'),62        t.callExpression(63          t.memberExpression(64            t.identifier('Symbol'),65            t.identifier('for'),66            false,67            false,68          ),69          [t.stringLiteral('react_forget')],70        ),71        true,72        false,73      ),74      true,75    ),76  );77}7879/**80 * COPIED FROM babel-plugin-react-compiler/src/Entrypoint/BabelUtils.ts81 */82type ComponentDeclaration = t.FunctionDeclaration & {83  __componentDeclaration: boolean;84};8586type BabelFn =87  | NodePath<t.FunctionDeclaration>88  | NodePath<t.FunctionExpression>89  | NodePath<t.ArrowFunctionExpression>;9091export function isComponentDeclaration(92  node: t.FunctionDeclaration,93): node is ComponentDeclaration {94  return Object.prototype.hasOwnProperty.call(node, '__componentDeclaration');95}9697/*98 * Adapted from the ESLint rule at99 * https://github.com/facebook/react/blob/main/packages/eslint-plugin-react-hooks/src/RulesOfHooks.js#L90-L103100 */101function isComponentOrHookLike(102  node: NodePath<103    t.FunctionDeclaration | t.ArrowFunctionExpression | t.FunctionExpression104  >,105): boolean {106  const functionName = getFunctionName(node);107  // Check if the name is component or hook like:108  if (functionName !== null && isComponentName(functionName)) {109    return (110      // As an added check we also look for hook invocations or JSX111      callsHooksOrCreatesJsx(node) &&112      /*113       * and avoid helper functions that take more than one argument114       * helpers are _usually_ named with lowercase, but some code may115       * violate this rule116       */117      node.get('params').length <= 1118    );119  } else if (functionName !== null && isHook(functionName)) {120    // Hooks have hook invocations or JSX, but can take any # of arguments121    return callsHooksOrCreatesJsx(node);122  }123124  /*125   * Otherwise for function or arrow function expressions, check if they126   * appear as the argument to React.forwardRef() or React.memo():127   */128  if (node.isFunctionExpression() || node.isArrowFunctionExpression()) {129    if (isForwardRefCallback(node) || isMemoCallback(node)) {130      // As an added check we also look for hook invocations or JSX131      return callsHooksOrCreatesJsx(node);132    } else {133      return false;134    }135  }136  return false;137}138139function isHookName(s: string): boolean {140  return /^use[A-Z0-9]/.test(s);141}142143/*144 * We consider hooks to be a hook name identifier or a member expression145 * containing a hook name.146 */147148function isHook(path: NodePath<t.Expression | t.PrivateName>): boolean {149  if (path.isIdentifier()) {150    return isHookName(path.node.name);151  } else if (152    path.isMemberExpression() &&153    !path.node.computed &&154    isHook(path.get('property'))155  ) {156    const obj = path.get('object').node;157    const isPascalCaseNameSpace = /^[A-Z].*/;158    return obj.type === 'Identifier' && isPascalCaseNameSpace.test(obj.name);159  } else {160    return false;161  }162}163164/*165 * Checks if the node is a React component name. React component names must166 * always start with an uppercase letter.167 */168169function isComponentName(path: NodePath<t.Expression>): boolean {170  return path.isIdentifier() && /^[A-Z]/.test(path.node.name);171}172/*173 * Checks if the node is a callback argument of forwardRef. This render function174 * should follow the rules of hooks.175 */176177function isForwardRefCallback(path: NodePath<t.Expression>): boolean {178  return !!(179    path.parentPath.isCallExpression() &&180    path.parentPath.get('callee').isExpression() &&181    isReactAPI(path.parentPath.get('callee'), 'forwardRef')182  );183}184185/*186 * Checks if the node is a callback argument of React.memo. This anonymous187 * functional component should follow the rules of hooks.188 */189190function isMemoCallback(path: NodePath<t.Expression>): boolean {191  return (192    path.parentPath.isCallExpression() &&193    path.parentPath.get('callee').isExpression() &&194    isReactAPI(path.parentPath.get('callee'), 'memo')195  );196}197198function isReactAPI(199  path: NodePath<t.Expression | t.PrivateName | t.V8IntrinsicIdentifier>,200  functionName: string,201): boolean {202  const node = path.node;203  return (204    (node.type === 'Identifier' && node.name === functionName) ||205    (node.type === 'MemberExpression' &&206      node.object.type === 'Identifier' &&207      node.object.name === 'React' &&208      node.property.type === 'Identifier' &&209      node.property.name === functionName)210  );211}212213function callsHooksOrCreatesJsx(node: NodePath<t.Node>): boolean {214  let invokesHooks = false;215  let createsJsx = false;216  node.traverse({217    JSX() {218      createsJsx = true;219    },220    CallExpression(call) {221      const callee = call.get('callee');222      if (callee.isExpression() && isHook(callee)) {223        invokesHooks = true;224      }225    },226  });227228  return invokesHooks || createsJsx;229}230231/*232 * Gets the static name of a function AST node. For function declarations it is233 * easy. For anonymous function expressions it is much harder. If you search for234 * `IsAnonymousFunctionDefinition()` in the ECMAScript spec you'll find places235 * where JS gives anonymous function expressions names. We roughly detect the236 * same AST nodes with some exceptions to better fit our use case.237 */238239function getFunctionName(240  path: NodePath<241    t.FunctionDeclaration | t.ArrowFunctionExpression | t.FunctionExpression242  >,243): NodePath<t.Expression> | null {244  if (path.isFunctionDeclaration()) {245    const id = path.get('id');246    if (id.isIdentifier()) {247      return id;248    }249    return null;250  }251  let id: NodePath<t.LVal | t.Expression | t.PrivateName> | null = null;252  const parent = path.parentPath;253  if (parent.isVariableDeclarator() && parent.get('init').node === path.node) {254    // const useHook = () => {};255    id = parent.get('id');256  } else if (257    parent.isAssignmentExpression() &&258    parent.get('right').node === path.node &&259    parent.get('operator') === '='260  ) {261    // useHook = () => {};262    id = parent.get('left');263  } else if (264    parent.isProperty() &&265    parent.get('value').node === path.node &&266    !parent.get('computed') &&267    parent.get('key').isLVal()268  ) {269    /*270     * {useHook: () => {}}271     * {useHook() {}}272     */273    id = parent.get('key');274  } else if (275    parent.isAssignmentPattern() &&276    parent.get('right').node === path.node &&277    !parent.get('computed')278  ) {279    /*280     * const {useHook = () => {}} = {};281     * ({useHook = () => {}} = {});282     *283     * Kinda clowny, but we'd said we'd follow spec convention for284     * `IsAnonymousFunctionDefinition()` usage.285     */286    id = parent.get('left');287  }288  if (id !== null && (id.isIdentifier() || id.isMemberExpression())) {289    return id;290  } else {291    return null;292  }293}