[compiler] Collect temporaries and optional chains from inner functions

Recursively collect identifier / property loads and optional chains from inner functions. This PR is in preparation for #31200

Previously, we only did this in `collectHoistablePropertyLoads` to understand hoistable property loads from inner functions.
1. collectTemporariesSidemap
2. collectOptionalChainSidemap
3. collectHoistablePropertyLoads
    - ^ this recursively calls `collectTemporariesSidemap`, `collectOptionalChainSidemap`, and `collectOptionalChainSidemap` on inner functions
4. collectDependencies

Now, we have
1. collectTemporariesSidemap
    - recursively record identifiers in inner functions. Note that we track all temporaries in the same map as `IdentifierIds` are currently unique across functions
2. collectOptionalChainSidemap
    - recursively records optional chain sidemaps in inner functions
3. collectHoistablePropertyLoads
    - (unchanged, except to remove recursive collection of temporaries)
4. collectDependencies
    - unchanged: to be modified to recursively collect dependencies in next PR

'

Author: mofeiZ

compiler/packages/babel-plugin-react-compiler/src/HIR/CollectHoistablePropertyLoads.ts

@@ -8,7 +8,6 @@ import {
   Set_union,
   getOrInsertDefault,
 } from '../Utils/utils';
-import {collectOptionalChainSidemap} from './CollectOptionalChainDependencies';
 import {
   BasicBlock,
   BlockId,
@@ -22,7 +21,6 @@ import {
   ReactiveScopeDependency,
   ScopeId,
 } from './HIR';
-import {collectTemporariesSidemap} from './PropagateScopeDependenciesHIR';
 
 const DEBUG_PRINT = false;
 
@@ -373,17 +371,10 @@ function collectNonNullsInBlocks(
         !fn.env.config.enableTreatFunctionDepsAsConditional
       ) {
         const innerFn = instr.value.loweredFunc;
-        const innerTemporaries = collectTemporariesSidemap(
-          innerFn.func,
-          new Set(),
-        );
-        const innerOptionals = collectOptionalChainSidemap(innerFn.func);
         const innerHoistableMap = collectHoistablePropertyLoadsImpl(
           innerFn.func,
           {
             ...context,
-            temporaries: innerTemporaries, // TODO: remove in later PR
-            hoistableFromOptionals: innerOptionals.hoistableObjects, // TODO: remove in later PR
             nestedFnImmutableContext:
               context.nestedFnImmutableContext ??
               new Set(

compiler/packages/babel-plugin-react-compiler/src/HIR/CollectOptionalChainDependencies.ts

@@ -1,4 +1,5 @@
 import {CompilerError} from '..';
+import {getOrInsertDefault} from '../Utils/utils';
 import {assertNonNull} from './CollectHoistablePropertyLoads';
 import {
   BlockId,
@@ -22,25 +23,14 @@ export function collectOptionalChainSidemap(
   fn: HIRFunction,
 ): OptionalChainSidemap {
   const context: OptionalTraversalContext = {
+    currFn: fn,
     blocks: fn.body.blocks,
     seenOptionals: new Set(),
-    processedInstrsInOptional: new Set(),
+    processedInstrsInOptional: new Map(),
     temporariesReadInOptional: new Map(),
     hoistableObjects: new Map(),
   };
-  for (const [_, block] of fn.body.blocks) {
-    if (
-      block.terminal.kind === 'optional' &&
-      !context.seenOptionals.has(block.id)
-    ) {
-      traverseOptionalBlock(
-        block as TBasicBlock<OptionalTerminal>,
-        context,
-        null,
-      );
-    }
-  }
-
+  traverseFunction(fn, context);
   return {
     temporariesReadInOptional: context.temporariesReadInOptional,
     processedInstrsInOptional: context.processedInstrsInOptional,
@@ -96,8 +86,13 @@ export type OptionalChainSidemap = {
    * bb5:
    *   $5 = MethodCall $2.$4() <--- here, we want to take a dep on $2 and $4!
    * ```
+   *
+   * Also note that InstructionIds are not unique across inner functions.
    */
-  processedInstrsInOptional: ReadonlySet<InstructionId>;
+  processedInstrsInOptional: ReadonlyMap<
+    HIRFunction,
+    ReadonlySet<InstructionId>
+  >;
   /**
    * Records optional chains for which we can safely evaluate non-optional
    * PropertyLoads. e.g. given `a?.b.c`, we can evaluate any load from `a?.b` at
@@ -115,16 +110,46 @@ export type OptionalChainSidemap = {
 };
 
 type OptionalTraversalContext = {
+  currFn: HIRFunction;
   blocks: ReadonlyMap<BlockId, BasicBlock>;
 
   // Track optional blocks to avoid outer calls into nested optionals
   seenOptionals: Set<BlockId>;
 
-  processedInstrsInOptional: Set<InstructionId>;
+  processedInstrsInOptional: Map<HIRFunction, Set<InstructionId>>;
   temporariesReadInOptional: Map<IdentifierId, ReactiveScopeDependency>;
   hoistableObjects: Map<BlockId, ReactiveScopeDependency>;
 };
 
+function traverseFunction(
+  fn: HIRFunction,
+  context: OptionalTraversalContext,
+): void {
+  for (const [_, block] of fn.body.blocks) {
+    for (const instr of block.instructions) {
+      if (
+        instr.value.kind === 'FunctionExpression' ||
+        instr.value.kind === 'ObjectMethod'
+      ) {
+        traverseFunction(instr.value.loweredFunc.func, {
+          ...context,
+          currFn: instr.value.loweredFunc.func,
+          blocks: instr.value.loweredFunc.func.body.blocks,
+        });
+      }
+    }
+    if (
+      block.terminal.kind === 'optional' &&
+      !context.seenOptionals.has(block.id)
+    ) {
+      traverseOptionalBlock(
+        block as TBasicBlock<OptionalTerminal>,
+        context,
+        null,
+      );
+    }
+  }
+}
 /**
  * Match the consequent and alternate blocks of an optional.
  * @returns propertyload computed by the consequent block, or null if the
@@ -369,10 +394,13 @@ function traverseOptionalBlock(
       },
     ],
   };
-  context.processedInstrsInOptional.add(
-    matchConsequentResult.storeLocalInstrId,
+  const processedInstrsInOptionalByFn = getOrInsertDefault(
+    context.processedInstrsInOptional,
+    context.currFn,
+    new Set(),
   );
-  context.processedInstrsInOptional.add(test.id);
+  processedInstrsInOptionalByFn.add(matchConsequentResult.storeLocalInstrId);
+  processedInstrsInOptionalByFn.add(test.id);
   context.temporariesReadInOptional.set(
     matchConsequentResult.consequentId,
     load,

compiler/packages/babel-plugin-react-compiler/src/HIR/PropagateScopeDependenciesHIR.ts

@@ -176,8 +176,10 @@ function findTemporariesUsedOutsideDeclaringScope(
  * $2 = LoadLocal 'foo'
  * $3 = CallExpression $2($1)
  * ```
- * Only map LoadLocal and PropertyLoad lvalues to their source if we know that
- * reordering the read (from the time-of-load to time-of-use) is valid.
+ * @param usedOutsideDeclaringScope is used to check the correctness of
+ * reordering LoadLocal / PropertyLoad calls. We only track a LoadLocal /
+ * PropertyLoad in the returned temporaries map if reordering the read (from the
+ * time-of-load to time-of-use) is valid.
  *
  * If a LoadLocal or PropertyLoad instruction is within the reactive scope range
  * (a proxy for mutable range) of the load source, later instructions may
@@ -215,7 +217,29 @@ export function collectTemporariesSidemap(
   fn: HIRFunction,
   usedOutsideDeclaringScope: ReadonlySet<DeclarationId>,
 ): ReadonlyMap<IdentifierId, ReactiveScopeDependency> {
-  const temporaries = new Map<IdentifierId, ReactiveScopeDependency>();
+  const temporaries = new Map();
+  collectTemporariesSidemapImpl(
+    fn,
+    usedOutsideDeclaringScope,
+    temporaries,
+    false,
+  );
+  return temporaries;
+}
+
+/**
+ * Recursive collect a sidemap of all `LoadLocal` and `PropertyLoads` with a
+ * function and all nested functions.
+ *
+ * Note that IdentifierIds are currently unique, so we can use a single
+ * Map<IdentifierId, ...> across all nested functions.
+ */
+function collectTemporariesSidemapImpl(
+  fn: HIRFunction,
+  usedOutsideDeclaringScope: ReadonlySet<DeclarationId>,
+  temporaries: Map<IdentifierId, ReactiveScopeDependency>,
+  isInnerFn: boolean,
+): void {
   for (const [_, block] of fn.body.blocks) {
     for (const instr of block.instructions) {
       const {value, lvalue} = instr;
@@ -224,27 +248,51 @@ export function collectTemporariesSidemap(
       );
 
       if (value.kind === 'PropertyLoad' && !usedOutside) {
-        const property = getProperty(
-          value.object,
-          value.property,
-          false,
-          temporaries,
-        );
-        temporaries.set(lvalue.identifier.id, property);
+        if (!isInnerFn || temporaries.has(value.object.identifier.id)) {
+          /**
+           * All dependencies of a inner / nested function must have a base
+           * identifier from the outermost component / hook. This is because the
+           * compiler cannot break an inner function into multiple granular
+           * scopes.
+           */
+          const property = getProperty(
+            value.object,
+            value.property,
+            false,
+            temporaries,
+          );
+          temporaries.set(lvalue.identifier.id, property);
+        }
       } else if (
         value.kind === 'LoadLocal' &&
         lvalue.identifier.name == null &&
         value.place.identifier.name !== null &&
         !usedOutside
       ) {
-        temporaries.set(lvalue.identifier.id, {
-          identifier: value.place.identifier,
-          path: [],
-        });
+        if (
+          !isInnerFn ||
+          fn.context.some(
+            context => context.identifier.id === value.place.identifier.id,
+          )
+        ) {
+          temporaries.set(lvalue.identifier.id, {
+            identifier: value.place.identifier,
+            path: [],
+          });
+        }
+      } else if (
+        value.kind === 'FunctionExpression' ||
+        value.kind === 'ObjectMethod'
+      ) {
+        collectTemporariesSidemapImpl(
+          value.loweredFunc.func,
+          usedOutsideDeclaringScope,
+          temporaries,
+          true,
+        );
       }
     }
   }
-  return temporaries;
 }
 
 function getProperty(
@@ -310,6 +358,12 @@ class Context {
   #temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>;
   #temporariesUsedOutsideScope: ReadonlySet<DeclarationId>;
 
+  /**
+   * Tracks the traversal state. See Context.declare for explanation of why this
+   * is needed.
+   */
+  inInnerFn: boolean = false;
+
   constructor(
     temporariesUsedOutsideScope: ReadonlySet<DeclarationId>,
     temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
@@ -360,12 +414,23 @@ class Context {
   }
 
   /*
-   * Records where a value was declared, and optionally, the scope where the value originated from.
-   * This is later used to determine if a dependency should be added to a scope; if the current
-   * scope we are visiting is the same scope where the value originates, it can't be a dependency
-   * on itself.
+   * Records where a value was declared, and optionally, the scope where the
+   * value originated from. This is later used to determine if a dependency
+   * should be added to a scope; if the current scope we are visiting is the
+   * same scope where the value originates, it can't be a dependency on itself.
+   *
+   * Note that we do not track declarations or reassignments within inner
+   * functions for the following reasons:
+   *   - inner functions cannot be split by scope boundaries and are guaranteed
+   *     to consume their own declarations
+   *   - reassignments within inner functions are tracked as context variables,
+   *     which already have extended mutable ranges to account for reassignments
+   *   - *most importantly* it's currently simply incorrect to compare inner
+   *     function instruction ids (tracked by `decl`) with outer ones (as stored
+   *     by root identifier mutable ranges).
    */
   declare(identifier: Identifier, decl: Decl): void {
+    if (this.inInnerFn) return;
     if (!this.#declarations.has(identifier.declarationId)) {
       this.#declarations.set(identifier.declarationId, decl);
     }
@@ -575,7 +640,10 @@ function collectDependencies(
   fn: HIRFunction,
   usedOutsideDeclaringScope: ReadonlySet<DeclarationId>,
   temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
-  processedInstrsInOptional: ReadonlySet<InstructionId>,
+  processedInstrsInOptional: ReadonlyMap<
+    HIRFunction,
+    ReadonlySet<InstructionId>
+  >,
 ): Map<ReactiveScope, Array<ReactiveScopeDependency>> {
   const context = new Context(usedOutsideDeclaringScope, temporaries);
 
@@ -595,6 +663,12 @@ function collectDependencies(
 
   const scopeTraversal = new ScopeBlockTraversal();
 
+  const shouldSkipInstructionDependencies = (
+    fn: HIRFunction,
+    id: InstructionId,
+  ): boolean => {
+    return processedInstrsInOptional.get(fn)?.has(id) ?? false;
+  };
   for (const [blockId, block] of fn.body.blocks) {
     scopeTraversal.recordScopes(block);
     const scopeBlockInfo = scopeTraversal.blockInfos.get(blockId);
@@ -614,12 +688,12 @@ function collectDependencies(
       }
     }
     for (const instr of block.instructions) {
-      if (!processedInstrsInOptional.has(instr.id)) {
+      if (!shouldSkipInstructionDependencies(fn, instr.id)) {
         handleInstruction(instr, context);
       }
     }
 
-    if (!processedInstrsInOptional.has(block.terminal.id)) {
+    if (!shouldSkipInstructionDependencies(fn, block.terminal.id)) {
       for (const place of eachTerminalOperand(block.terminal)) {
         context.visitOperand(place);
       }

compiler/packages/babel-plugin-react-compiler/src/HIR/visitors.ts

@@ -1215,9 +1215,17 @@ export class ScopeBlockTraversal {
     }
   }
 
+  /**
+   * @returns if the given scope is currently 'active', i.e. if the scope start
+   * block but not the scope fallthrough has been recorded.
+   */
   isScopeActive(scopeId: ScopeId): boolean {
     return this.#activeScopes.indexOf(scopeId) !== -1;
   }
+
+  /**
+   * The current, innermost active scope.
+   */
   get currentScope(): ScopeId | null {
     return this.#activeScopes.at(-1) ?? null;
   }