PruneScopesTransform Class — react Architecture

Source Code

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts lines 1011–1123

class PruneScopesTransform extends ReactiveFunctionTransform<
  Set<DeclarationId>
> {
  prunedScopes: Set<ScopeId> = new Set();
  /**
   * Track reassignments so we can correctly set `pruned` flags for
   * inlined useMemos.
   */
  reassignments: Map<DeclarationId, Set<Identifier>> = new Map();

  override transformScope(
    scopeBlock: ReactiveScopeBlock,
    state: Set<DeclarationId>,
  ): Transformed<ReactiveStatement> {
    this.visitScope(scopeBlock, state);

    /**
     * Scopes may initially appear "empty" because the value being memoized
     * is early-returned from within the scope. For now we intentionaly keep
     * these scopes, and let them get pruned later by PruneUnusedScopes
     * _after_ handling the early-return case in PropagateEarlyReturns.
     *
     * Also keep the scope if an early return was created by some earlier pass,
     * which may happen in alternate compiler configurations.
     */
    if (
      (scopeBlock.scope.declarations.size === 0 &&
        scopeBlock.scope.reassignments.size === 0) ||
      scopeBlock.scope.earlyReturnValue !== null
    ) {
      return {kind: 'keep'};
    }

    const hasMemoizedOutput =
      Array.from(scopeBlock.scope.declarations.values()).some(decl =>
        state.has(decl.identifier.declarationId),
      ) ||
      Array.from(scopeBlock.scope.reassignments).some(identifier =>
        state.has(identifier.declarationId),
      );
    if (hasMemoizedOutput) {
      return {kind: 'keep'};
    } else {
      this.prunedScopes.add(scopeBlock.scope.id);
      return {
        kind: 'replace-many',
        value: scopeBlock.instructions,
      };
    }
  }

  /**
   * If we pruned the scope for a non-escaping value, we know it doesn't
   * need to be memoized. Remove associated `Memoize` instructions so that
   * we don't report false positives on "missing" memoization of these values.
   */
  override transformInstruction(
    instruction: ReactiveInstruction,
    state: Set<DeclarationId>,
  ): Transformed<ReactiveStatement> {
    this.traverseInstruction(instruction, state);

    const value = instruction.value;
    if (value.kind === 'StoreLocal' && value.lvalue.kind === 'Reassign') {
      // Complex cases of useMemo inlining result in a temporary that is reassigned
      const ids = getOrInsertDefault(
        this.reassignments,
        value.lvalue.place.identifier.declarationId,
        new Set(),
      );
      ids.add(value.value.identifier);
    } else if (
      value.kind === 'LoadLocal' &&
      value.place.identifier.scope != null &&
      instruction.lvalue != null &&
      instruction.lvalue.identifier.scope == null
    ) {
      /*
       * Simpler cases result in a direct assignment to the original lvalue, with a
       * LoadLocal
       */