[compiler] First cut at dep inference (#31386)

This is for researching/prototyping, not a feature we are releasing
imminently.

Putting up an early version of inferring effect dependencies to get
feedback on the approach. We do not plan to ship this as-is, and may not
start by going after direct `useEffect` calls. Until we make that
decision, the heuristic I use to detect when to insert effect deps will
suffice for testing.

The approach is simple: when we see a useEffect call with no dep array
we insert the deps inferred for the lambda passed in. If the first
argument is not a lambda then we do not do anything.

This diff is the easy part. I think the harder part will be ensuring
that we can infer the deps even when we have to bail out of memoization.
We have no other features that *must* run regardless of rules of react
violations. Does anyone foresee any issues using the compiler passes to
infer reactive deps when there may be violations?

I have a few questions:
1. Will there ever be more than one instruction in a block containing a
useEffect? if no, I can get rid of the`addedInstrs` variable that I use
to make sure I insert the effect deps array temp creation at the right
spot.
2. Are there any cases for resolving the first argument beyond just
looking at the lvalue's identifier id that I'll need to take into
account? e.g., do I need to recursively resolve certain bindings?

---------

Co-authored-by: Mofei Zhang <[email protected]>

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

@@ -36,6 +36,7 @@ import {
   inferReactivePlaces,
   inferReferenceEffects,
   inlineImmediatelyInvokedFunctionExpressions,
+  inferEffectDependencies,
 } from '../Inference';
 import {
   constantPropagation,
@@ -354,6 +355,10 @@ function* runWithEnvironment(
     value: hir,
   });
 
+  if (env.config.inferEffectDependencies) {
+    inferEffectDependencies(env, hir);
+  }
+
   if (env.config.inlineJsxTransform) {
     inlineJsxTransform(hir, env.config.inlineJsxTransform);
     yield log({

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

@@ -233,6 +233,19 @@ const EnvironmentConfigSchema = z.object({
 
   enableFunctionDependencyRewrite: z.boolean().default(true),
 
+  /**
+   * 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),
+
+  /**
+   * Enables inference and auto-insertion of effect dependencies. Still experimental.
+   */
+  inferEffectDependencies: z.boolean().default(false),
+
   /**
    * Enables inlining ReactElement object literals in place of JSX
    * An alternative to the standard JSX transform which replaces JSX with React's jsxProd() runtime

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

@@ -0,0 +1,247 @@
+import {CompilerError, SourceLocation} from '..';
+import {
+  ArrayExpression,
+  Effect,
+  Environment,
+  FunctionExpression,
+  GeneratedSource,
+  HIRFunction,
+  IdentifierId,
+  Instruction,
+  isUseEffectHookType,
+  makeInstructionId,
+  TInstruction,
+  InstructionId,
+  ScopeId,
+  ReactiveScopeDependency,
+  Place,
+  ReactiveScopeDependencies,
+} from '../HIR';
+import {
+  createTemporaryPlace,
+  fixScopeAndIdentifierRanges,
+  markInstructionIds,
+} from '../HIR/HIRBuilder';
+import {eachInstructionOperand, eachTerminalOperand} from '../HIR/visitors';
+
+/**
+ * Infers reactive dependencies captured by useEffect lambdas and adds them as
+ * a second argument to the useEffect call if no dependency array is provided.
+ */
+export function inferEffectDependencies(
+  env: Environment,
+  fn: HIRFunction,
+): void {
+  let hasRewrite = false;
+  const fnExpressions = new Map<
+    IdentifierId,
+    TInstruction<FunctionExpression>
+  >();
+  const scopeInfos = new Map<
+    ScopeId,
+    {pruned: boolean; deps: ReactiveScopeDependencies; hasSingleInstr: boolean}
+  >();
+
+  /**
+   * When inserting LoadLocals, we need to retain the reactivity of the base
+   * identifier, as later passes e.g. PruneNonReactiveDeps take the reactivity of
+   * a base identifier as the "maximal" reactivity of all its references.
+   * Concretely,
+   * reactive(Identifier i) = Union_{reference of i}(reactive(reference))
+   */
+  const reactiveIds = inferReactiveIdentifiers(fn);
+
+  for (const [, block] of fn.body.blocks) {
+    if (
+      block.terminal.kind === 'scope' ||
+      block.terminal.kind === 'pruned-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,
+      });
+    }
+    const rewriteInstrs = new Map<InstructionId, Array<Instruction>>();
+    for (const instr of block.instructions) {
+      const {value, lvalue} = instr;
+      if (value.kind === 'FunctionExpression') {
+        fnExpressions.set(
+          lvalue.identifier.id,
+          instr as TInstruction<FunctionExpression>,
+        );
+      } else if (
+        /*
+         * This check is not final. Right now we only look for useEffects without a dependency array.
+         * This is likely not how we will ship this feature, but it is good enough for us to make progress
+         * on the implementation and test it.
+         */
+        value.kind === 'CallExpression' &&
+        isUseEffectHookType(value.callee.identifier) &&
+        value.args.length === 1 &&
+        value.args[0].kind === 'Identifier'
+      ) {
+        const fnExpr = fnExpressions.get(value.args[0].identifier.id);
+        if (fnExpr != null) {
+          const scopeInfo =
+            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,
+            });
+          }
+
+          /**
+           * Step 1: write new instructions to insert a dependency array
+           *
+           * Note that it's invalid to prune non-reactive deps in this pass, see
+           * the `infer-effect-deps/pruned-nonreactive-obj` fixture for an
+           * explanation.
+           */
+          const effectDeps: Array<Place> = [];
+          const newInstructions: Array<Instruction> = [];
+          for (const dep of scopeInfo.deps) {
+            const {place, instructions} = writeDependencyToInstructions(
+              dep,
+              reactiveIds.has(dep.identifier.id),
+              fn.env,
+              fnExpr.loc,
+            );
+            newInstructions.push(...instructions);
+            effectDeps.push(place);
+          }
+          const deps: ArrayExpression = {
+            kind: 'ArrayExpression',
+            elements: effectDeps,
+            loc: GeneratedSource,
+          };
+
+          const depsPlace = createTemporaryPlace(env, GeneratedSource);
+          depsPlace.effect = Effect.Read;
+
+          newInstructions.push({
+            id: makeInstructionId(0),
+            loc: GeneratedSource,
+            lvalue: {...depsPlace, effect: Effect.Mutate},
+            value: deps,
+          });
+
+          // Step 2: insert the deps array as an argument of the useEffect
+          value.args[1] = {...depsPlace, effect: Effect.Freeze};
+          rewriteInstrs.set(instr.id, newInstructions);
+        }
+      }
+    }
+    if (rewriteInstrs.size > 0) {
+      hasRewrite = true;
+      const newInstrs = [];
+      for (const instr of block.instructions) {
+        const newInstr = rewriteInstrs.get(instr.id);
+        if (newInstr != null) {
+          newInstrs.push(...newInstr, instr);
+        } else {
+          newInstrs.push(instr);
+        }
+      }
+      block.instructions = newInstrs;
+    }
+  }
+  if (hasRewrite) {
+    // Renumber instructions and fix scope ranges
+    markInstructionIds(fn.body);
+    fixScopeAndIdentifierRanges(fn.body);
+  }
+}
+
+function writeDependencyToInstructions(
+  dep: ReactiveScopeDependency,
+  reactive: boolean,
+  env: Environment,
+  loc: SourceLocation,
+): {place: Place; instructions: Array<Instruction>} {
+  const instructions: Array<Instruction> = [];
+  let currValue = createTemporaryPlace(env, GeneratedSource);
+  currValue.reactive = reactive;
+  instructions.push({
+    id: makeInstructionId(0),
+    loc: GeneratedSource,
+    lvalue: {...currValue, effect: Effect.Mutate},
+    value: {
+      kind: 'LoadLocal',
+      place: {
+        kind: 'Identifier',
+        identifier: dep.identifier,
+        effect: Effect.Capture,
+        reactive,
+        loc: loc,
+      },
+      loc: loc,
+    },
+  });
+  for (const path of dep.path) {
+    if (path.optional) {
+      /**
+       * TODO: instead of truncating optional paths, reuse
+       * instructions from hoisted dependencies block(s)
+       */
+      break;
+    }
+    const nextValue = createTemporaryPlace(env, GeneratedSource);
+    nextValue.reactive = reactive;
+    instructions.push({
+      id: makeInstructionId(0),
+      loc: GeneratedSource,
+      lvalue: {...nextValue, effect: Effect.Mutate},
+      value: {
+        kind: 'PropertyLoad',
+        object: {...currValue, effect: Effect.Capture},
+        property: path.property,
+        loc: loc,
+      },
+    });
+    currValue = nextValue;
+  }
+  currValue.effect = Effect.Freeze;
+  return {place: currValue, instructions};
+}
+
+function inferReactiveIdentifiers(fn: HIRFunction): Set<IdentifierId> {
+  const reactiveIds: Set<IdentifierId> = new Set();
+  for (const [, block] of fn.body.blocks) {
+    for (const instr of block.instructions) {
+      /**
+       * No need to traverse into nested functions as
+       * 1. their effects are recorded in `LoweredFunction.dependencies`
+       * 2. we don't mark `reactive` in these anyways
+       */
+      for (const place of eachInstructionOperand(instr)) {
+        if (place.reactive) {
+          reactiveIds.add(place.identifier.id);
+        }
+      }
+    }
+
+    for (const place of eachTerminalOperand(block.terminal)) {
+      if (place.reactive) {
+        reactiveIds.add(place.identifier.id);
+      }
+    }
+  }
+  return reactiveIds;
+}

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

@@ -11,3 +11,4 @@ export {inferMutableRanges} from './InferMutableRanges';
 export {inferReactivePlaces} from './InferReactivePlaces';
 export {default as inferReferenceEffects} from './InferReferenceEffects';
 export {inlineImmediatelyInvokedFunctionExpressions} from './InlineImmediatelyInvokedFunctionExpressions';
+export {inferEffectDependencies} from './InferEffectDependencies';