[compiler] Stop using function dependencies in propagateScopeDeps

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts

@@ -57,7 +57,6 @@ import {
   mergeReactiveScopesThatInvalidateTogether,
   promoteUsedTemporaries,
   propagateEarlyReturns,
-  propagateScopeDependencies,
   pruneHoistedContexts,
   pruneNonEscapingScopes,
   pruneNonReactiveDependencies,
@@ -348,14 +347,12 @@ function* runWithEnvironment(
   });
   assertTerminalSuccessorsExist(hir);
   assertTerminalPredsExist(hir);
-  if (env.config.enablePropagateDepsInHIR) {
-    propagateScopeDependenciesHIR(hir);
-    yield log({
-      kind: 'hir',
-      name: 'PropagateScopeDependenciesHIR',
-      value: hir,
-    });
-  }
+  propagateScopeDependenciesHIR(hir);
+  yield log({
+    kind: 'hir',
+    name: 'PropagateScopeDependenciesHIR',
+    value: hir,
+  });
 
   if (env.config.inlineJsxTransform) {
     inlineJsxTransform(hir, env.config.inlineJsxTransform);
@@ -383,15 +380,6 @@ function* runWithEnvironment(
   });
   assertScopeInstructionsWithinScopes(reactiveFunction);
 
-  if (!env.config.enablePropagateDepsInHIR) {
-    propagateScopeDependencies(reactiveFunction);
-    yield log({
-      kind: 'reactive',
-      name: 'PropagateScopeDependencies',
-      value: reactiveFunction,
-    });
-  }
-
   pruneNonEscapingScopes(reactiveFunction);
   yield log({
     kind: 'reactive',

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

@@ -1,13 +1,13 @@
 import {CompilerError} from '../CompilerError';
 import {inRange} from '../ReactiveScopes/InferReactiveScopeVariables';
+import {printDependency} from '../ReactiveScopes/PrintReactiveFunction';
 import {
   Set_equal,
   Set_filter,
   Set_intersect,
   Set_union,
   getOrInsertDefault,
 } from '../Utils/utils';
-import {collectOptionalChainSidemap} from './CollectOptionalChainDependencies';
 import {
   BasicBlock,
   BlockId,
@@ -21,7 +21,8 @@ import {
   ReactiveScopeDependency,
   ScopeId,
 } from './HIR';
-import {collectTemporariesSidemap} from './PropagateScopeDependenciesHIR';
+
+const DEBUG_PRINT = false;
 
 /**
  * Helper function for `PropagateScopeDependencies`. Uses control flow graph
@@ -86,15 +87,8 @@ export function collectHoistablePropertyLoads(
   fn: HIRFunction,
   temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
   hoistableFromOptionals: ReadonlyMap<BlockId, ReactiveScopeDependency>,
-  nestedFnImmutableContext: ReadonlySet<IdentifierId> | null,
 ): ReadonlyMap<BlockId, BlockInfo> {
   const registry = new PropertyPathRegistry();
-
-  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
@@ -111,14 +105,51 @@ export function collectHoistablePropertyLoads(
       }
     }
   }
-  const nodes = collectNonNullsInBlocks(fn, {
-    temporaries: actuallyEvaluatedTemporaries,
+  return collectHoistablePropertyLoadsImpl(fn, {
+    temporaries,
     knownImmutableIdentifiers,
     hoistableFromOptionals,
     registry,
-    nestedFnImmutableContext,
+    nestedFnImmutableContext: null,
   });
-  propagateNonNull(fn, nodes, registry);
+}
+
+type CollectHoistablePropertyLoadsContext = {
+  temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>;
+  knownImmutableIdentifiers: ReadonlySet<IdentifierId>;
+  hoistableFromOptionals: ReadonlyMap<BlockId, ReactiveScopeDependency>;
+  registry: PropertyPathRegistry;
+  /**
+   * (For nested / inner function declarations)
+   * Context variables (i.e. captured from an outer scope) that are immutable.
+   * Note that this technically could be merged into `knownImmutableIdentifiers`,
+   * but are currently kept separate for readability.
+   */
+  nestedFnImmutableContext: ReadonlySet<IdentifierId> | null;
+};
+function collectHoistablePropertyLoadsImpl(
+  fn: HIRFunction,
+  context: CollectHoistablePropertyLoadsContext,
+): ReadonlyMap<BlockId, BlockInfo> {
+  const functionExpressionLoads = collectFunctionExpressionFakeLoads(fn);
+  const actuallyEvaluatedTemporaries = new Map(
+    [...context.temporaries].filter(([id]) => !functionExpressionLoads.has(id)),
+  );
+
+  const nodes = collectNonNullsInBlocks(fn, {
+    ...context,
+    temporaries: actuallyEvaluatedTemporaries,
+  });
+  propagateNonNull(fn, nodes, context.registry);
+
+  if (DEBUG_PRINT) {
+    console.log('(printing hoistable nodes in blocks)');
+    for (const [blockId, node] of nodes) {
+      console.log(
+        `bb${blockId}: ${[...node.assumedNonNullObjects].map(n => printDependency(n.fullPath)).join(' ')}`,
+      );
+    }
+  }
 
   return nodes;
 }
@@ -243,7 +274,7 @@ class PropertyPathRegistry {
 
 function getMaybeNonNullInInstruction(
   instr: InstructionValue,
-  context: CollectNonNullsInBlocksContext,
+  context: CollectHoistablePropertyLoadsContext,
 ): PropertyPathNode | null {
   let path = null;
   if (instr.kind === 'PropertyLoad') {
@@ -262,7 +293,7 @@ function getMaybeNonNullInInstruction(
 function isImmutableAtInstr(
   identifier: Identifier,
   instr: InstructionId,
-  context: CollectNonNullsInBlocksContext,
+  context: CollectHoistablePropertyLoadsContext,
 ): boolean {
   if (context.nestedFnImmutableContext != null) {
     /**
@@ -295,22 +326,9 @@ function isImmutableAtInstr(
   }
 }
 
-type CollectNonNullsInBlocksContext = {
-  temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>;
-  knownImmutableIdentifiers: ReadonlySet<IdentifierId>;
-  hoistableFromOptionals: ReadonlyMap<BlockId, ReactiveScopeDependency>;
-  registry: PropertyPathRegistry;
-  /**
-   * (For nested / inner function declarations)
-   * Context variables (i.e. captured from an outer scope) that are immutable.
-   * Note that this technically could be merged into `knownImmutableIdentifiers`,
-   * but are currently kept separate for readability.
-   */
-  nestedFnImmutableContext: ReadonlySet<IdentifierId> | null;
-};
 function collectNonNullsInBlocks(
   fn: HIRFunction,
-  context: CollectNonNullsInBlocksContext,
+  context: CollectHoistablePropertyLoadsContext,
 ): ReadonlyMap<BlockId, BlockInfo> {
   /**
    * Known non-null objects such as functional component props can be safely
@@ -348,27 +366,25 @@ function collectNonNullsInBlocks(
         assumedNonNullObjects.add(maybeNonNull);
       }
       if (
-        instr.value.kind === 'FunctionExpression' &&
+        (instr.value.kind === 'FunctionExpression' ||
+          instr.value.kind === 'ObjectMethod') &&
         !fn.env.config.enableTreatFunctionDepsAsConditional
       ) {
         const innerFn = instr.value.loweredFunc;
-        const innerTemporaries = collectTemporariesSidemap(
-          innerFn.func,
-          new Set(),
-        );
-        const innerOptionals = collectOptionalChainSidemap(innerFn.func);
-        const innerHoistableMap = collectHoistablePropertyLoads(
+        const innerHoistableMap = collectHoistablePropertyLoadsImpl(
           innerFn.func,
-          innerTemporaries,
-          innerOptionals.hoistableObjects,
-          context.nestedFnImmutableContext ??
-            new Set(
-              innerFn.func.context
-                .filter(place =>
-                  isImmutableAtInstr(place.identifier, instr.id, context),
-                )
-                .map(place => place.identifier.id),
-            ),
+          {
+            ...context,
+            nestedFnImmutableContext:
+              context.nestedFnImmutableContext ??
+              new Set(
+                innerFn.func.context
+                  .filter(place =>
+                    isImmutableAtInstr(place.identifier, instr.id, context),
+                  )
+                  .map(place => place.identifier.id),
+              ),
+          },
         );
         const innerHoistables = assertNonNull(
           innerHoistableMap.get(innerFn.func.body.entry),
@@ -591,7 +607,10 @@ function collectFunctionExpressionFakeLoads(
 
   for (const [_, block] of fn.body.blocks) {
     for (const {lvalue, value} of block.instructions) {
-      if (value.kind === 'FunctionExpression') {
+      if (
+        value.kind === 'FunctionExpression' ||
+        value.kind === 'ObjectMethod'
+      ) {
         for (const reference of value.loweredFunc.dependencies) {
           let curr: IdentifierId | undefined = reference.identifier.id;
           while (curr != null) {

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/Environment.ts

@@ -230,15 +230,7 @@ const EnvironmentConfigSchema = z.object({
    */
   enableUseTypeAnnotations: z.boolean().default(false),
 
-  enablePropagateDepsInHIR: z.boolean().default(false),
-
-  /**
-   * Enables inference of optional dependency chains. Without this flag
-   * a property chain such as `props?.items?.foo` will infer as a dep on
-   * just `props`. With this flag enabled, we'll infer that full path as
-   * the dependency.
-   */
-  enableOptionalDependencies: z.boolean().default(true),
+  enableFunctionDependencyRewrite: z.boolean().default(true),
 
   /**
    * Enables inlining ReactElement object literals in place of JSX