lowerContextAccess() — react Function Reference

Source Code

compiler/packages/babel-plugin-react-compiler/src/Optimization/LowerContextAccess.ts lines 36–132

export function lowerContextAccess(
  fn: HIRFunction,
  loweredContextCalleeConfig: ExternalFunction,
): void {
  const contextAccess: Map<IdentifierId, CallExpression> = new Map();
  const contextKeys: Map<IdentifierId, Array<string>> = new Map();

  // collect context access and keys
  for (const [, block] of fn.body.blocks) {
    for (const instr of block.instructions) {
      const {value, lvalue} = instr;

      if (
        value.kind === 'CallExpression' &&
        isUseContextHookType(value.callee.identifier)
      ) {
        contextAccess.set(lvalue.identifier.id, value);
        continue;
      }

      if (value.kind !== 'Destructure') {
        continue;
      }

      const destructureId = value.value.identifier.id;
      if (!contextAccess.has(destructureId)) {
        continue;
      }

      const keys = getContextKeys(value);
      if (keys === null) {
        return;
      }

      if (contextKeys.has(destructureId)) {
        /*
         * TODO(gsn): Add support for accessing context over multiple
         * statements.
         */
        return;
      } else {
        contextKeys.set(destructureId, keys);
      }
    }
  }

  let importLoweredContextCallee: NonLocalImportSpecifier | null = null;

  if (contextAccess.size > 0 && contextKeys.size > 0) {
    for (const [, block] of fn.body.blocks) {
      let nextInstructions: Array<Instruction> | null = null;

      for (let i = 0; i < block.instructions.length; i++) {
        const instr = block.instructions[i];
        const {lvalue, value} = instr;
        if (
          value.kind === 'CallExpression' &&
          isUseContextHookType(value.callee.identifier) &&
          contextKeys.has(lvalue.identifier.id)
        ) {
          importLoweredContextCallee ??=
            fn.env.programContext.addImportSpecifier(
              loweredContextCalleeConfig,
            );
          const loweredContextCalleeInstr = emitLoadLoweredContextCallee(
            fn.env,
            importLoweredContextCallee,
          );

          if (nextInstructions === null) {
            nextInstructions = block.instructions.slice(0, i);
          }
          nextInstructions.push(loweredContextCalleeInstr);

          const keys = contextKeys.get(lvalue.identifier.id)!;
          const selectorFnInstr = emitSelectorFn(fn.env, keys);
          nextInstructions.push(selectorFnInstr);

          const lowerContextCallId = loweredContextCalleeInstr.lvalue;
          value.callee = lowerContextCallId;