insertAdditionalFunctionDeclaration() — react Function Reference

Source Code

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Gating.ts lines 36–126

function insertAdditionalFunctionDeclaration(
  fnPath: NodePath<t.FunctionDeclaration>,
  compiled: t.FunctionDeclaration,
  programContext: ProgramContext,
  gatingFunctionIdentifierName: string,
): void {
  const originalFnName = fnPath.node.id;
  const originalFnParams = fnPath.node.params;
  const compiledParams = fnPath.node.params;
  /**
   * Note that other than `export default function() {}`, all other function
   * declarations must have a binding identifier. Since default exports cannot
   * be referenced, it's safe to assume that all function declarations passed
   * here will have an identifier.
   * https://tc39.es/ecma262/multipage/ecmascript-language-functions-and-classes.html#sec-function-definitions
   */
  CompilerError.invariant(originalFnName != null && compiled.id != null, {
    reason:
      'Expected function declarations that are referenced elsewhere to have a named identifier',
    loc: fnPath.node.loc ?? GeneratedSource,
  });
  CompilerError.invariant(originalFnParams.length === compiledParams.length, {
    reason:
      'Expected React Compiler optimized function declarations to have the same number of parameters as source',
    loc: fnPath.node.loc ?? GeneratedSource,
  });

  const gatingCondition = t.identifier(
    programContext.newUid(`${gatingFunctionIdentifierName}_result`),
  );
  const unoptimizedFnName = t.identifier(
    programContext.newUid(`${originalFnName.name}_unoptimized`),
  );
  const optimizedFnName = t.identifier(
    programContext.newUid(`${originalFnName.name}_optimized`),
  );
  /**
   * Step 1: rename existing functions
   */
  compiled.id.name = optimizedFnName.name;
  fnPath.get('id').replaceInline(unoptimizedFnName);

  /**
   * Step 2: insert new function declaration
   */
  const newParams: Array<t.Identifier | t.RestElement> = [];
  const genNewArgs: Array<() => t.Identifier | t.SpreadElement> = [];
  for (let i = 0; i < originalFnParams.length; i++) {
    const argName = `arg${i}`;
    if (originalFnParams[i].type === 'RestElement') {
      newParams.push(t.restElement(t.identifier(argName)));
      genNewArgs.push(() => t.spreadElement(t.identifier(argName)));
    } else {
      newParams.push(t.identifier(argName));
      genNewArgs.push(() => t.identifier(argName));
    }
  }
  // insertAfter called in reverse order of how nodes should appear in program
  fnPath.insertAfter(
    t.functionDeclaration(
      originalFnName,
      newParams,
      t.blockStatement([
        t.ifStatement(
          gatingCondition,
          t.returnStatement(
            t.callExpression(
              compiled.id,
              genNewArgs.map(fn => fn()),
            ),
          ),
          t.returnStatement(
            t.callExpression(
              unoptimizedFnName,
              genNewArgs.map(fn => fn()),
            ),
          ),
        ),
      ]),
    ),
  );