[compiler][hir] Only hoist always-accessed PropertyLoads from function decls

ghstack-source-id: 485a44f
Pull Request resolved: #31066

Author: mofeiZ

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

@@ -7,6 +7,7 @@ import {
   Set_union,
   getOrInsertDefault,
 } from '../Utils/utils';
+import {collectOptionalChainSidemap} from './CollectOptionalChainDependencies';
 import {
   BasicBlock,
   BlockId,
@@ -15,10 +16,13 @@ import {
   HIRFunction,
   Identifier,
   IdentifierId,
+  Instruction,
+  InstructionId,
   InstructionValue,
   ReactiveScopeDependency,
   ScopeId,
 } from './HIR';
+import {collectTemporariesSidemap} from './PropagateScopeDependenciesHIR';
 
 /**
  * Helper function for `PropagateScopeDependencies`.
@@ -73,30 +77,73 @@ export function collectHoistablePropertyLoads(
   fn: HIRFunction,
   temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
   optionals: ReadonlyMap<BlockId, ReactiveScopeDependency>,
-): ReadonlyMap<ScopeId, BlockInfo> {
+  nestedFnContext: NestedFunctionContext | null,
+): ReadonlyMap<BlockId, BlockInfo> {
   const registry = new PropertyPathRegistry();
 
-  const nodes = collectNonNullsInBlocks(fn, temporaries, optionals, registry);
+  const functionExpressionLoads = collectFunctionExpressionFakeLoads(fn);
+  const actuallyEvaluatedTemporaries = new Map(
+    [...temporaries].filter(([id]) => !functionExpressionLoads.has(id)),
+  );
+
+  /**
+   * Due to current limitations of mutable range inference, there are edge cases in
+   * which we infer known-immutable values (e.g. props or hook params) to have a
+   * mutable range and scope.
+   * (see `destructure-array-declaration-to-context-var` fixture)
+   * We track known immutable identifiers to reduce regressions (as PropagateScopeDeps
+   * is being rewritten to HIR).
+   */
+  const knownImmutableIdentifiers = new Set<IdentifierId>();
+  if (fn.fnType === 'Component' || fn.fnType === 'Hook') {
+    for (const p of fn.params) {
+      if (p.kind === 'Identifier') {
+        knownImmutableIdentifiers.add(p.identifier.id);
+      }
+    }
+  }
+  const nodes = collectNonNullsInBlocks(fn, {
+    temporaries: actuallyEvaluatedTemporaries,
+    knownImmutableIdentifiers,
+    optionals,
+    registry,
+    nestedFnContext,
+  });
   propagateNonNull(fn, nodes, registry);
 
-  const nodesKeyedByScopeId = new Map<ScopeId, BlockInfo>();
+  return nodes;
+}
+
+export function keyByScopeId<T>(
+  fn: HIRFunction,
+  source: ReadonlyMap<BlockId, T>,
+): ReadonlyMap<ScopeId, T> {
+  const keyedByScopeId = new Map<ScopeId, T>();
   for (const [_, block] of fn.body.blocks) {
     if (block.terminal.kind === 'scope') {
-      nodesKeyedByScopeId.set(
+      keyedByScopeId.set(
         block.terminal.scope.id,
-        nodes.get(block.terminal.block)!,
+        source.get(block.terminal.block)!,
       );
     }
   }
-
-  return nodesKeyedByScopeId;
+  return keyedByScopeId;
 }
 
 export type BlockInfo = {
   block: BasicBlock;
   assumedNonNullObjects: ReadonlySet<PropertyPathNode>;
 };
 
+type NestedFunctionContext = {
+  /**
+   * Instruction id of the outermost nested function declaration.
+   */
+  outermostFnInstrId: InstructionId;
+  // Captured from outermost scope
+  capturedIdentifiers: Set<IdentifierId>;
+};
+
 /**
  * PropertyLoadRegistry data structure to dedupe property loads (e.g. a.b.c)
  * and make computing sets intersections simpler.
@@ -196,45 +243,33 @@ class PropertyPathRegistry {
 
 function getMaybeNonNullInInstruction(
   instr: InstructionValue,
-  temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
-  registry: PropertyPathRegistry,
+  context: CollectNonNullsInBlocksContext,
 ): PropertyPathNode | null {
   let path = null;
   if (instr.kind === 'PropertyLoad') {
-    path = temporaries.get(instr.object.identifier.id) ?? {
+    path = context.temporaries.get(instr.object.identifier.id) ?? {
       identifier: instr.object.identifier,
       path: [],
     };
   } else if (instr.kind === 'Destructure') {
-    path = temporaries.get(instr.value.identifier.id) ?? null;
+    path = context.temporaries.get(instr.value.identifier.id) ?? null;
   } else if (instr.kind === 'ComputedLoad') {
-    path = temporaries.get(instr.object.identifier.id) ?? null;
+    path = context.temporaries.get(instr.object.identifier.id) ?? null;
   }
-  return path != null ? registry.getOrCreateProperty(path) : null;
+  return path != null ? context.registry.getOrCreateProperty(path) : null;
 }
 
+type CollectNonNullsInBlocksContext = {
+  temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>;
+  knownImmutableIdentifiers: ReadonlySet<IdentifierId>;
+  optionals: ReadonlyMap<BlockId, ReactiveScopeDependency>;
+  registry: PropertyPathRegistry;
+  nestedFnContext: NestedFunctionContext | null;
+};
 function collectNonNullsInBlocks(
   fn: HIRFunction,
-  temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
-  optionals: ReadonlyMap<BlockId, ReactiveScopeDependency>,
-  registry: PropertyPathRegistry,
+  context: CollectNonNullsInBlocksContext,
 ): ReadonlyMap<BlockId, BlockInfo> {
-  /**
-   * Due to current limitations of mutable range inference, there are edge cases in
-   * which we infer known-immutable values (e.g. props or hook params) to have a
-   * mutable range and scope.
-   * (see `destructure-array-declaration-to-context-var` fixture)
-   * We track known immutable identifiers to reduce regressions (as PropagateScopeDeps
-   * is being rewritten to HIR).
-   */
-  const knownImmutableIdentifiers = new Set<IdentifierId>();
-  if (fn.fnType === 'Component' || fn.fnType === 'Hook') {
-    for (const p of fn.params) {
-      if (p.kind === 'Identifier') {
-        knownImmutableIdentifiers.add(p.identifier.id);
-      }
-    }
-  }
   /**
    * Known non-null objects such as functional component props can be safely
    * read from any block.
@@ -246,7 +281,9 @@ function collectNonNullsInBlocks(
     fn.params[0].kind === 'Identifier'
   ) {
     const identifier = fn.params[0].identifier;
-    knownNonNullIdentifiers.add(registry.getOrCreateIdentifier(identifier));
+    knownNonNullIdentifiers.add(
+      context.registry.getOrCreateIdentifier(identifier),
+    );
   }
   const nodes = new Map<BlockId, BlockInfo>();
   for (const [_, block] of fn.body.blocks) {
@@ -258,20 +295,22 @@ function collectNonNullsInBlocks(
       block,
       assumedNonNullObjects,
     });
-    const maybeOptionalChain = optionals.get(block.id);
+    const maybeOptionalChain = context.optionals.get(block.id);
     if (maybeOptionalChain != null) {
       assumedNonNullObjects.add(
-        registry.getOrCreateProperty(maybeOptionalChain),
+        context.registry.getOrCreateProperty(maybeOptionalChain),
       );
       continue;
     }
     for (const instr of block.instructions) {
-      const maybeNonNull = getMaybeNonNullInInstruction(
-        instr.value,
-        temporaries,
-        registry,
-      );
-      if (maybeNonNull != null) {
+      const maybeNonNull = getMaybeNonNullInInstruction(instr.value, context);
+      if (
+        maybeNonNull != null &&
+        (context.nestedFnContext == null ||
+          context.nestedFnContext.capturedIdentifiers.has(
+            maybeNonNull.fullPath.identifier.id,
+          ))
+      ) {
         const baseIdentifier = maybeNonNull.fullPath.identifier;
         /**
          * Since this runs *after* buildReactiveScopeTerminals, identifier mutable ranges
@@ -289,17 +328,48 @@ function collectNonNullsInBlocks(
           baseIdentifier.scope != null &&
           inRange(
             {
-              id: instr.id,
+              id:
+                context.nestedFnContext != null
+                  ? context.nestedFnContext.outermostFnInstrId
+                  : instr.id,
             },
             baseIdentifier.scope.range,
           );
         if (
           !isMutableAtInstr ||
-          knownImmutableIdentifiers.has(baseIdentifier.id)
+          context.knownImmutableIdentifiers.has(baseIdentifier.id)
         ) {
           assumedNonNullObjects.add(maybeNonNull);
         }
       }
+      if (
+        instr.value.kind === 'FunctionExpression' &&
+        !fn.env.config.enableTreatFunctionDepsAsConditional
+      ) {
+        const innerFn = instr.value.loweredFunc;
+        const innerTemporaries = collectTemporariesSidemap(
+          innerFn.func,
+          new Set(),
+        );
+        const innerOptionals = collectOptionalChainSidemap(innerFn.func);
+        const innerHoistableMap = collectHoistablePropertyLoads(
+          innerFn.func,
+          innerTemporaries,
+          innerOptionals.hoistableObjects,
+          context.nestedFnContext ?? {
+            outermostFnInstrId: instr.id,
+            capturedIdentifiers: new Set(
+              innerFn.func.context.map(place => place.identifier.id),
+            ),
+          },
+        );
+        const innerHoistables = assertNonNull(
+          innerHoistableMap.get(innerFn.func.body.entry),
+        );
+        for (const entry of innerHoistables.assumedNonNullObjects) {
+          assumedNonNullObjects.add(entry);
+        }
+      }
     }
   }
   return nodes;
@@ -500,3 +570,27 @@ function reduceMaybeOptionalChains(
     }
   } while (changed);
 }
+
+function collectFunctionExpressionFakeLoads(
+  fn: HIRFunction,
+): Set<IdentifierId> {
+  const sources = new Map<IdentifierId, IdentifierId>();
+  const functionExpressionReferences = new Set<IdentifierId>();
+
+  for (const [_, block] of fn.body.blocks) {
+    for (const {lvalue, value} of block.instructions) {
+      if (value.kind === 'FunctionExpression') {
+        for (const reference of value.loweredFunc.dependencies) {
+          let curr: IdentifierId | undefined = reference.identifier.id;
+          while (curr != null) {
+            functionExpressionReferences.add(curr);
+            curr = sources.get(curr);
+          }
+        }
+      } else if (value.kind === 'PropertyLoad') {
+        sources.set(lvalue.identifier.id, value.object.identifier.id);
+      }
+    }
+  }
+  return functionExpressionReferences;
+}

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

@@ -17,7 +17,10 @@ import {
   areEqualPaths,
   IdentifierId,
 } from './HIR';
-import {collectHoistablePropertyLoads} from './CollectHoistablePropertyLoads';
+import {
+  collectHoistablePropertyLoads,
+  keyByScopeId,
+} from './CollectHoistablePropertyLoads';
 import {
   ScopeBlockTraversal,
   eachInstructionOperand,
@@ -41,10 +44,9 @@ export function propagateScopeDependenciesHIR(fn: HIRFunction): void {
     hoistableObjects,
   } = collectOptionalChainSidemap(fn);
 
-  const hoistablePropertyLoads = collectHoistablePropertyLoads(
+  const hoistablePropertyLoads = keyByScopeId(
     fn,
-    temporaries,
-    hoistableObjects,
+    collectHoistablePropertyLoads(fn, temporaries, hoistableObjects, null),
   );
 
   const scopeDeps = collectDependencies(
@@ -209,7 +211,7 @@ function findTemporariesUsedOutsideDeclaringScope(
  * of $1, as the evaluation of `arr.length` changes between instructions $1 and
  * $3. We do not track $1 -> arr.length in this case.
  */
-function collectTemporariesSidemap(
+export function collectTemporariesSidemap(
   fn: HIRFunction,
   usedOutsideDeclaringScope: ReadonlySet<DeclarationId>,
 ): ReadonlyMap<IdentifierId, ReactiveScopeDependency> {