react/compiler/packages/babel-plugin-react-compiler/src/Inference/InlineImmediatelyInvokedFunctionExpressions.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 {
  BasicBlock,
  BlockId,
  Effect,
  Environment,
  FunctionExpression,
  GeneratedSource,
  GotoVariant,
  HIRFunction,
  IdentifierId,
  InstructionKind,
  LabelTerminal,
  Place,
  makeInstructionId,
  makeType,
  promoteTemporary,
  reversePostorderBlocks,
} from "../HIR";
import { markInstructionIds, markPredecessors } from "../HIR/HIRBuilder";
import { eachInstructionValueOperand } from "../HIR/visitors";
import { retainWhere } from "../Utils/utils";

/*
 * Inlines immediately invoked function expressions (IIFEs) to allow more fine-grained memoization
 * of the values they produce.
 *
 * Example:
 *
 * ```
 * const x = (() => {
 *    const x = [];
 *    x.push(foo());
 *    return x;
 * })();
 *
 * =>
 *
 * bb0:
 *     // placeholder for the result, all return statements will assign here
 *    let t0;
 *    // Label allows using a goto (break) to exit out of the body
 *    Label block=bb1 fallthrough=bb2
 * bb1:
 *    // code within the function expression
 *    const x0 = [];
 *    x0.push(foo());
 *    // return is replaced by assignment to the result variable...
 *    t0 = x0;
 *    // ...and a goto to the code after the function expression invocation
 *    Goto bb2
 * bb2:
 *    // code after the IIFE call
 *    const x = t0;
 * ```
 *
 * The implementation relies on HIR's ability to support labeled blocks:
 * - We terminate the basic block just prior to the CallExpression of the IIFE
 *    with a LabelTerminal whose fallback is the code following the CallExpression.
 *    Just prior to the terminal we also create a named temporary variable which
 *    will hold the result.
 * - We then inline the contents of the function "in between" (conceptually) those
 *    two blocks.
 * - All return statements in the original function expression are replaced with a
 *    StoreLocal to the temporary we allocated before plus a Goto to the fallthrough
 *    block (code following the CallExpression).
 */
export function inlineImmediatelyInvokedFunctionExpressions(
  fn: HIRFunction
): void {
  // Track all function expressions that are assigned to a temporary
  const functions = new Map<IdentifierId, FunctionExpression>();
  // Functions that are inlined
  const inlinedFunctions = new Set<IdentifierId>();

  /*
   * Iterate the *existing* blocks from the outer component to find IIFEs
   * and inline them. During iteration we will modify `fn` (by inlining the CFG
   * of IIFEs) so we explicitly copy references to just the original
   * function's blocks first. As blocks are split to make room for IIFE calls,
   * the split portions of the blocks will be added to this queue.
   */
  const queue = Array.from(fn.body.blocks.values());
  queue: for (const block of queue) {
    for (let ii = 0; ii < block.instructions.length; ii++) {
      const instr = block.instructions[ii]!;
      switch (instr.value.kind) {
        case "FunctionExpression": {
          if (instr.lvalue.identifier.name === null) {
            functions.set(instr.lvalue.identifier.id, instr.value);
          }
          break;
        }
        case "CallExpression": {
          if (instr.value.args.length !== 0) {
            // We don't support inlining when there are arguments
            continue;
          }
          const body = functions.get(instr.value.callee.identifier.id);
          if (body === undefined) {
            // Not invoking a local function expression, can't inline
            continue;
          }

          if (
            body.loweredFunc.func.params.length > 0 ||
            body.loweredFunc.func.async ||
            body.loweredFunc.func.generator
          ) {
            // Can't inline functions with params, or async/generator functions
            continue;
          }

          // We know this function is used for an IIFE and can prune it later
          inlinedFunctions.add(instr.value.callee.identifier.id);

          // Create a new block which will contain code following the IIFE call
          const continuationBlockId = fn.env.nextBlockId;
          const continuationBlock: BasicBlock = {
            id: continuationBlockId,
            instructions: block.instructions.slice(ii + 1),
            kind: block.kind,
            phis: new Set(),
            preds: new Set(),
            terminal: block.terminal,
          };
          fn.body.blocks.set(continuationBlockId, continuationBlock);

          /*
           * Trim the original block to contain instructions up to (but not including)
           * the IIFE
           */
          block.instructions.length = ii;

          /*
           * To account for complex control flow within the lambda, we treat the lambda
           * as if it were a single labeled statement, and replace all returns with gotos
           * to the label fallthrough.
           */
          const newTerminal: LabelTerminal = {
            block: body.loweredFunc.func.body.entry,
            id: makeInstructionId(0),
            kind: "label",
            fallthrough: continuationBlockId,
            loc: block.terminal.loc,
          };
          block.terminal = newTerminal;

          // We store the result in the IIFE temporary
          const result = instr.lvalue;

          // Declare the IIFE temporary
          declareTemporary(fn.env, block, result);

          // Promote the temporary with a name as we require this to persist
          promoteTemporary(result.identifier);

          /*
           * Rewrite blocks from the lambda to replace any `return` with a
           * store to the result and `goto` the continuation block
           */
          for (const [id, block] of body.loweredFunc.func.body.blocks) {
            block.preds.clear();
            rewriteBlock(fn.env, block, continuationBlockId, result);
            fn.body.blocks.set(id, block);
          }

          /*
           * Ensure we visit the continuation block, since there may have been
           * sequential IIFEs that need to be visited.
           */
          queue.push(continuationBlock);
          continue queue;
        }
        default: {
          for (const place of eachInstructionValueOperand(instr.value)) {
            // Any other use of a function expression means it isn't an IIFE
            functions.delete(place.identifier.id);
          }
        }
      }
    }
  }

  if (inlinedFunctions.size !== 0) {
    // Remove instructions that define lambdas which we inlined
    for (const [, block] of fn.body.blocks) {
      retainWhere(
        block.instructions,
        (instr) => !inlinedFunctions.has(instr.lvalue.identifier.id)
      );
    }

    /*
     * If terminals have changed then blocks may have become newly unreachable.
     * Re-run minification of the graph (incl reordering instruction ids)
     */
    reversePostorderBlocks(fn.body);
    markInstructionIds(fn.body);
    markPredecessors(fn.body);
  }
}

/*
 * Rewrites the block so that all `return` terminals are replaced:
 * * Add a StoreLocal <returnValue> = <terminal.value>
 * * Replace the terminal with a Goto to <returnTarget>
 */
function rewriteBlock(
  env: Environment,
  block: BasicBlock,
  returnTarget: BlockId,
  returnValue: Place
): void {
  const { terminal } = block;
  if (terminal.kind !== "return") {
    return;
  }
  block.instructions.push({
    id: makeInstructionId(0),
    loc: terminal.loc,
    lvalue: {
      effect: Effect.Unknown,
      identifier: {
        id: env.nextIdentifierId,
        mutableRange: {
          start: makeInstructionId(0),
          end: makeInstructionId(0),
        },
        name: null,
        scope: null,
        type: makeType(),
        loc: terminal.loc,
      },
      kind: "Identifier",
      reactive: false,
      loc: terminal.loc,
    },
    value: {
      kind: "StoreLocal",
      lvalue: { kind: InstructionKind.Reassign, place: { ...returnValue } },
      value: terminal.value,
      type: null,
      loc: terminal.loc,
    },
  });
  block.terminal = {
    kind: "goto",
    block: returnTarget,
    id: makeInstructionId(0),
    variant: GotoVariant.Break,
    loc: block.terminal.loc,
  };
}

function declareTemporary(
  env: Environment,
  block: BasicBlock,
  result: Place
): void {
  block.instructions.push({
    id: makeInstructionId(0),
    loc: GeneratedSource,
    lvalue: {
      effect: Effect.Unknown,
      identifier: {
        id: env.nextIdentifierId,
        mutableRange: {
          start: makeInstructionId(0),
          end: makeInstructionId(0),
        },
        name: null,
        scope: null,
        type: makeType(),
        loc: result.loc,
      },
      kind: "Identifier",
      reactive: false,
      loc: GeneratedSource,
    },
    value: {
      kind: "DeclareLocal",
      lvalue: {
        place: result,
        kind: InstructionKind.Let,
      },
      type: null,
      loc: result.loc,
    },
  });
}