[compiler] Fallback for inferred effect dependencies (facebook#32984)

When effect dependencies cannot be inferred due to memoization-related
bailouts or unexpected mutable ranges (which currently often have to do
with writes to refs), fall back to traversing the effect lambda itself.

This fallback uses the same logic as PropagateScopeDependencies:
1. Collect a sidemap of loads and property loads
2. Find hoistable accesses from the control flow graph. Note that here,
we currently take into account the mutable ranges of instructions (see
`mutate-after-useeffect-granular-access` fixture)
3. Collect the set of property paths accessed by the effect
4. Merge to get the set of minimal dependencies

Author: mofeiZ

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

@@ -19,6 +19,7 @@ import {
   BasicBlock,
   BlockId,
   DependencyPathEntry,
+  FunctionExpression,
   GeneratedSource,
   getHookKind,
   HIRFunction,
@@ -30,6 +31,7 @@ import {
   PropertyLiteral,
   ReactiveScopeDependency,
   ScopeId,
+  TInstruction,
 } from './HIR';
 
 const DEBUG_PRINT = false;
@@ -127,6 +129,33 @@ export function collectHoistablePropertyLoads(
   });
 }
 
+export function collectHoistablePropertyLoadsInInnerFn(
+  fnInstr: TInstruction<FunctionExpression>,
+  temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
+  hoistableFromOptionals: ReadonlyMap<BlockId, ReactiveScopeDependency>,
+): ReadonlyMap<BlockId, BlockInfo> {
+  const fn = fnInstr.value.loweredFunc.func;
+  const initialContext: CollectHoistablePropertyLoadsContext = {
+    temporaries,
+    knownImmutableIdentifiers: new Set(),
+    hoistableFromOptionals,
+    registry: new PropertyPathRegistry(),
+    nestedFnImmutableContext: null,
+    assumedInvokedFns: fn.env.config.enableTreatFunctionDepsAsConditional
+      ? new Set()
+      : getAssumedInvokedFunctions(fn),
+  };
+  const nestedFnImmutableContext = new Set(
+    fn.context
+      .filter(place =>
+        isImmutableAtInstr(place.identifier, fnInstr.id, initialContext),
+      )
+      .map(place => place.identifier.id),
+  );
+  initialContext.nestedFnImmutableContext = nestedFnImmutableContext;
+  return collectHoistablePropertyLoadsImpl(fn, initialContext);
+}
+
 type CollectHoistablePropertyLoadsContext = {
   temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>;
   knownImmutableIdentifiers: ReadonlySet<IdentifierId>;

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

@@ -116,7 +116,7 @@ export function propagateScopeDependenciesHIR(fn: HIRFunction): void {
   }
 }
 
-function findTemporariesUsedOutsideDeclaringScope(
+export function findTemporariesUsedOutsideDeclaringScope(
   fn: HIRFunction,
 ): ReadonlySet<DeclarationId> {
   /*
@@ -378,7 +378,7 @@ type Decl = {
   scope: Stack<ReactiveScope>;
 };
 
-class Context {
+export class DependencyCollectionContext {
   #declarations: Map<DeclarationId, Decl> = new Map();
   #reassignments: Map<Identifier, Decl> = new Map();
 
@@ -645,7 +645,10 @@ enum HIRValue {
   Terminal,
 }
 
-function handleInstruction(instr: Instruction, context: Context): void {
+export function handleInstruction(
+  instr: Instruction,
+  context: DependencyCollectionContext,
+): void {
   const {id, value, lvalue} = instr;
   context.declare(lvalue.identifier, {
     id,
@@ -708,7 +711,7 @@ function collectDependencies(
   temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
   processedInstrsInOptional: ReadonlySet<Instruction | Terminal>,
 ): Map<ReactiveScope, Array<ReactiveScopeDependency>> {
-  const context = new Context(
+  const context = new DependencyCollectionContext(
     usedOutsideDeclaringScope,
     temporaries,
     processedInstrsInOptional,

compiler/packages/babel-plugin-react-compiler/src/Inference/InferEffectDependencies.ts

@@ -22,18 +22,30 @@ import {
   ScopeId,
   ReactiveScopeDependency,
   Place,
+  ReactiveScope,
   ReactiveScopeDependencies,
+  Terminal,
   isUseRefType,
   isSetStateType,
   isFireFunctionType,
+  makeScopeId,
 } from '../HIR';
+import {collectHoistablePropertyLoadsInInnerFn} from '../HIR/CollectHoistablePropertyLoads';
+import {collectOptionalChainSidemap} from '../HIR/CollectOptionalChainDependencies';
+import {ReactiveScopeDependencyTreeHIR} from '../HIR/DeriveMinimalDependenciesHIR';
 import {DEFAULT_EXPORT} from '../HIR/Environment';
 import {
   createTemporaryPlace,
   fixScopeAndIdentifierRanges,
   markInstructionIds,
 } from '../HIR/HIRBuilder';
+import {
+  collectTemporariesSidemap,
+  DependencyCollectionContext,
+  handleInstruction,
+} from '../HIR/PropagateScopeDependenciesHIR';
 import {eachInstructionOperand, eachTerminalOperand} from '../HIR/visitors';
+import {empty} from '../Utils/Stack';
 import {getOrInsertWith} from '../Utils/utils';
 
 /**
@@ -62,10 +74,7 @@ export function inferEffectDependencies(fn: HIRFunction): void {
   const autodepFnLoads = new Map<IdentifierId, number>();
   const autodepModuleLoads = new Map<IdentifierId, Map<string, number>>();
 
-  const scopeInfos = new Map<
-    ScopeId,
-    {pruned: boolean; deps: ReactiveScopeDependencies; hasSingleInstr: boolean}
-  >();
+  const scopeInfos = new Map<ScopeId, ReactiveScopeDependencies>();
 
   const loadGlobals = new Set<IdentifierId>();
 
@@ -79,19 +88,18 @@ export function inferEffectDependencies(fn: HIRFunction): void {
   const reactiveIds = inferReactiveIdentifiers(fn);
 
   for (const [, block] of fn.body.blocks) {
-    if (
-      block.terminal.kind === 'scope' ||
-      block.terminal.kind === 'pruned-scope'
-    ) {
+    if (block.terminal.kind === 'scope') {
       const scopeBlock = fn.body.blocks.get(block.terminal.block)!;
-      scopeInfos.set(block.terminal.scope.id, {
-        pruned: block.terminal.kind === 'pruned-scope',
-        deps: block.terminal.scope.dependencies,
-        hasSingleInstr:
-          scopeBlock.instructions.length === 1 &&
-          scopeBlock.terminal.kind === 'goto' &&
-          scopeBlock.terminal.block === block.terminal.fallthrough,
-      });
+      if (
+        scopeBlock.instructions.length === 1 &&
+        scopeBlock.terminal.kind === 'goto' &&
+        scopeBlock.terminal.block === block.terminal.fallthrough
+      ) {
+        scopeInfos.set(
+          block.terminal.scope.id,
+          block.terminal.scope.dependencies,
+        );
+      }
     }
     const rewriteInstrs = new Map<InstructionId, Array<Instruction>>();
     for (const instr of block.instructions) {
@@ -173,31 +181,22 @@ export function inferEffectDependencies(fn: HIRFunction): void {
               fnExpr.lvalue.identifier.scope != null
                 ? scopeInfos.get(fnExpr.lvalue.identifier.scope.id)
                 : null;
-            CompilerError.invariant(scopeInfo != null, {
-              reason: 'Expected function expression scope to exist',
-              loc: value.loc,
-            });
-            if (scopeInfo.pruned || !scopeInfo.hasSingleInstr) {
-              /**
-               * TODO: retry pipeline that ensures effect function expressions
-               * are placed into their own scope
-               */
-              CompilerError.throwTodo({
-                reason:
-                  '[InferEffectDependencies] Expected effect function to have non-pruned scope and its scope to have exactly one instruction',
-                loc: fnExpr.loc,
-              });
+            let minimalDeps: Set<ReactiveScopeDependency>;
+            if (scopeInfo != null) {
+              minimalDeps = new Set(scopeInfo);
+            } else {
+              minimalDeps = inferMinimalDependencies(fnExpr);
             }
-
             /**
              * Step 1: push dependencies to the effect deps array
              *
              * Note that it's invalid to prune all non-reactive deps in this pass, see
              * the `infer-effect-deps/pruned-nonreactive-obj` fixture for an
              * explanation.
              */
+
             const usedDeps = [];
-            for (const dep of scopeInfo.deps) {
+            for (const dep of minimalDeps) {
               if (
                 ((isUseRefType(dep.identifier) ||
                   isSetStateType(dep.identifier)) &&
@@ -422,3 +421,132 @@ function collectDepUsages(
 
   return sourceLocations;
 }
+
+function inferMinimalDependencies(
+  fnInstr: TInstruction<FunctionExpression>,
+): Set<ReactiveScopeDependency> {
+  const fn = fnInstr.value.loweredFunc.func;
+
+  const temporaries = collectTemporariesSidemap(fn, new Set());
+  const {
+    hoistableObjects,
+    processedInstrsInOptional,
+    temporariesReadInOptional,
+  } = collectOptionalChainSidemap(fn);
+
+  const hoistablePropertyLoads = collectHoistablePropertyLoadsInInnerFn(
+    fnInstr,
+    temporaries,
+    hoistableObjects,
+  );
+  const hoistableToFnEntry = hoistablePropertyLoads.get(fn.body.entry);
+  CompilerError.invariant(hoistableToFnEntry != null, {
+    reason:
+      '[InferEffectDependencies] Internal invariant broken: missing entry block',
+    loc: fnInstr.loc,
+  });
+
+  const dependencies = inferDependencies(
+    fnInstr,
+    new Map([...temporaries, ...temporariesReadInOptional]),
+    processedInstrsInOptional,
+  );
+
+  const tree = new ReactiveScopeDependencyTreeHIR(
+    [...hoistableToFnEntry.assumedNonNullObjects].map(o => o.fullPath),
+  );
+  for (const dep of dependencies) {
+    tree.addDependency({...dep});
+  }
+
+  return tree.deriveMinimalDependencies();
+}
+
+function inferDependencies(
+  fnInstr: TInstruction<FunctionExpression>,
+  temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
+  processedInstrsInOptional: ReadonlySet<Instruction | Terminal>,
+): Set<ReactiveScopeDependency> {
+  const fn = fnInstr.value.loweredFunc.func;
+  const context = new DependencyCollectionContext(
+    new Set(),
+    temporaries,
+    processedInstrsInOptional,
+  );
+  for (const dep of fn.context) {
+    context.declare(dep.identifier, {
+      id: makeInstructionId(0),
+      scope: empty(),
+    });
+  }
+  const placeholderScope: ReactiveScope = {
+    id: makeScopeId(0),
+    range: {
+      start: fnInstr.id,
+      end: makeInstructionId(fnInstr.id + 1),
+    },
+    dependencies: new Set(),
+    reassignments: new Set(),
+    declarations: new Map(),
+    earlyReturnValue: null,
+    merged: new Set(),
+    loc: GeneratedSource,
+  };
+  context.enterScope(placeholderScope);
+  inferDependenciesInFn(fn, context, temporaries);
+  context.exitScope(placeholderScope, false);
+  const resultUnfiltered = context.deps.get(placeholderScope);
+  CompilerError.invariant(resultUnfiltered != null, {
+    reason:
+      '[InferEffectDependencies] Internal invariant broken: missing scope dependencies',
+    loc: fn.loc,
+  });
+
+  const fnContext = new Set(fn.context.map(dep => dep.identifier.id));
+  const result = new Set<ReactiveScopeDependency>();
+  for (const dep of resultUnfiltered) {
+    if (fnContext.has(dep.identifier.id)) {
+      result.add(dep);
+    }
+  }
+
+  return result;
+}
+
+function inferDependenciesInFn(
+  fn: HIRFunction,
+  context: DependencyCollectionContext,
+  temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
+): void {
+  for (const [, block] of fn.body.blocks) {
+    // Record referenced optional chains in phis
+    for (const phi of block.phis) {
+      for (const operand of phi.operands) {
+        const maybeOptionalChain = temporaries.get(operand[1].identifier.id);
+        if (maybeOptionalChain) {
+          context.visitDependency(maybeOptionalChain);
+        }
+      }
+    }
+    for (const instr of block.instructions) {
+      if (
+        instr.value.kind === 'FunctionExpression' ||
+        instr.value.kind === 'ObjectMethod'
+      ) {
+        context.declare(instr.lvalue.identifier, {
+          id: instr.id,
+          scope: context.currentScope,
+        });
+        /**
+         * Recursively visit the inner function to extract dependencies
+         */
+        const innerFn = instr.value.loweredFunc.func;
+        context.enterInnerFn(instr as TInstruction<FunctionExpression>, () => {
+          inferDependenciesInFn(innerFn, context, temporaries);
+        });
+      } else {
+        handleInstruction(instr, context);
+      }
+    }
+  }
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/bailout-retry/mutate-after-useeffect-granular-access.expect.md

@@ -0,0 +1,39 @@
+
+## Input
+
+```javascript
+// @inferEffectDependencies @panicThreshold(none)
+import {useEffect} from 'react';
+import {print} from 'shared-runtime';
+
+function Component({foo}) {
+  const arr = [];
+  // Taking either arr[0].value or arr as a dependency is reasonable
+  // as long as developers know what to expect.
+  useEffect(() => print(arr[0].value));
+  arr.push({value: foo});
+  return arr;
+}
+
+```
+
+## Code
+
+```javascript
+// @inferEffectDependencies @panicThreshold(none)
+import { useEffect } from "react";
+import { print } from "shared-runtime";
+
+function Component(t0) {
+  const { foo } = t0;
+  const arr = [];
+
+  useEffect(() => print(arr[0].value), [arr[0].value]);
+  arr.push({ value: foo });
+  return arr;
+}
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/bailout-retry/mutate-after-useeffect-granular-access.js

@@ -0,0 +1,12 @@
+// @inferEffectDependencies @panicThreshold(none)
+import {useEffect} from 'react';
+import {print} from 'shared-runtime';
+
+function Component({foo}) {
+  const arr = [];
+  // Taking either arr[0].value or arr as a dependency is reasonable
+  // as long as developers know what to expect.
+  useEffect(() => print(arr[0].value));
+  arr.push({value: foo});
+  return arr;
+}