[compiler] Inferred effect dependencies now include optional chains

Inferred effect dependencies now include optional chains.

This is a temporary solution while https://github.com/facebook/react/pull/32099 and its followups are worked on. Ideally, we should model reactive scope dependencies in the IR similarly to `ComputeIR` -- dependencies should be hoisted and all references rewritten to use the hoisted dependencies.

(will add more test cases)

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

@@ -0,0 +1,280 @@
+import {
+  Place,
+  ReactiveScopeDependency,
+  Identifier,
+  makeInstructionId,
+  InstructionKind,
+  GeneratedSource,
+  BlockId,
+  makeTemporaryIdentifier,
+  Effect,
+  GotoVariant,
+  HIR,
+} from './HIR';
+import {CompilerError} from '../CompilerError';
+import {Environment} from './Environment';
+import HIRBuilder from './HIRBuilder';
+import {lowerValueToTemporary} from './BuildHIR';
+import {printDependency} from '../ReactiveScopes/PrintReactiveFunction';
+
+type DependencyInstructions = {
+  place: Place;
+  value: HIR;
+  exitBlockId: BlockId;
+};
+
+export function buildDependencyInstructions(
+  dep: ReactiveScopeDependency,
+  env: Environment,
+): DependencyInstructions {
+  const builder = new HIRBuilder(env, {
+    entryBlockKind: 'value',
+  });
+  let result: Place;
+  if (dep.path.every(path => !path.optional)) {
+    const last = writeNonOptionalDependency(dep, env, builder);
+    result = {
+      kind: 'Identifier',
+      identifier: last,
+      effect: Effect.Freeze,
+      reactive: dep.reactive,
+      loc: GeneratedSource,
+    };
+  } else {
+    const last = writeOptionalDependency(
+      dep.path.length - 1,
+      dep,
+      builder,
+      null,
+    );
+    result = {
+      kind: 'Identifier',
+      identifier: last,
+      effect: Effect.Freeze,
+      reactive: dep.reactive,
+      loc: GeneratedSource,
+    };
+  }
+
+  const exitBlockId = builder.terminate(
+    {
+      kind: 'unsupported',
+      loc: GeneratedSource,
+      id: makeInstructionId(0),
+    },
+    null,
+  );
+  return {
+    place: result,
+    value: builder.build(),
+    exitBlockId,
+  };
+}
+function writeNonOptionalDependency(
+  dep: ReactiveScopeDependency,
+  env: Environment,
+  builder: HIRBuilder,
+): Identifier {
+  const loc = dep.identifier.loc;
+  let last: Identifier = makeTemporaryIdentifier(env.nextIdentifierId, loc);
+  builder.push({
+    lvalue: {
+      identifier: last,
+      kind: 'Identifier',
+      effect: Effect.Mutate,
+      reactive: dep.reactive,
+      loc,
+    },
+    value: {
+      kind: 'LoadLocal',
+      place: {
+        identifier: dep.identifier,
+        kind: 'Identifier',
+        effect: Effect.Freeze,
+        reactive: dep.reactive,
+        loc,
+      },
+      loc,
+    },
+    id: makeInstructionId(1),
+    loc: loc,
+  });
+
+  for (const path of dep.path) {
+    const next = makeTemporaryIdentifier(env.nextIdentifierId, loc);
+    builder.push({
+      lvalue: {
+        identifier: next,
+        kind: 'Identifier',
+        effect: Effect.Mutate,
+        reactive: dep.reactive,
+        loc,
+      },
+      value: {
+        kind: 'PropertyLoad',
+        object: {
+          identifier: last,
+          kind: 'Identifier',
+          effect: Effect.Freeze,
+          reactive: dep.reactive,
+          loc,
+        },
+        property: path.property,
+        loc,
+      },
+      id: makeInstructionId(1),
+      loc: loc,
+    });
+    last = next;
+  }
+  return last;
+}
+
+function writeOptionalDependency(
+  idx: number,
+  dep: ReactiveScopeDependency,
+  builder: HIRBuilder,
+  parentAlternate: BlockId | null,
+): Identifier {
+  const env = builder.environment;
+  CompilerError.invariant(
+    idx >= 0 && !dep.path.slice(0, idx + 1).every(path => !path.optional),
+    {
+      reason: '[WriteOptional] Expected optional path',
+      description: `${idx} ${printDependency(dep)}`,
+      loc: GeneratedSource,
+    },
+  );
+  const continuationBlock = builder.reserve(builder.currentBlockKind());
+  const consequent = builder.reserve('value');
+
+  const returnPlace: Place = {
+    kind: 'Identifier',
+    identifier: makeTemporaryIdentifier(env.nextIdentifierId, GeneratedSource),
+    effect: Effect.Mutate,
+    reactive: dep.reactive,
+    loc: GeneratedSource,
+  };
+
+  let alternate;
+  if (parentAlternate != null) {
+    alternate = parentAlternate;
+  } else {
+    /**
+     * Make outermost alternate block
+     * $N = Primitive undefined
+     * $M = StoreLocal $OptionalResult = $N
+     * goto fallthrough
+     */
+    alternate = builder.enter('value', () => {
+      const temp = lowerValueToTemporary(builder, {
+        kind: 'Primitive',
+        value: undefined,
+        loc: GeneratedSource,
+      });
+      lowerValueToTemporary(builder, {
+        kind: 'StoreLocal',
+        lvalue: {kind: InstructionKind.Const, place: {...returnPlace}},
+        value: {...temp},
+        type: null,
+        loc: GeneratedSource,
+      });
+      return {
+        kind: 'goto',
+        variant: GotoVariant.Break,
+        block: continuationBlock.id,
+        id: makeInstructionId(0),
+        loc: GeneratedSource,
+      };
+    });
+  }
+
+  let testIdentifier: Identifier | null = null;
+  const testBlock = builder.enter('value', () => {
+    const firstOptional = dep.path.findIndex(path => path.optional);
+    if (idx === firstOptional) {
+      // Lower test block
+      testIdentifier = writeNonOptionalDependency(
+        {
+          identifier: dep.identifier,
+          reactive: dep.reactive,
+          path: dep.path.slice(0, idx),
+        },
+        env,
+        builder,
+      );
+    } else {
+      testIdentifier = writeOptionalDependency(
+        idx - 1,
+        dep,
+        builder,
+        alternate,
+      );
+    }
+
+    return {
+      kind: 'branch',
+      test: {
+        identifier: testIdentifier,
+        effect: Effect.Freeze,
+        kind: 'Identifier',
+        loc: GeneratedSource,
+        reactive: dep.reactive,
+      },
+      consequent: consequent.id,
+      alternate,
+      id: makeInstructionId(0),
+      loc: GeneratedSource,
+      fallthrough: continuationBlock.id,
+    };
+  });
+
+  builder.enterReserved(consequent, () => {
+    CompilerError.invariant(testIdentifier !== null, {
+      reason: 'Satisfy type checker',
+      description: null,
+      loc: null,
+      suggestions: null,
+    });
+
+    const tmpConsequent = lowerValueToTemporary(builder, {
+      kind: 'PropertyLoad',
+      object: {
+        identifier: testIdentifier,
+        kind: 'Identifier',
+        effect: Effect.Freeze,
+        reactive: dep.reactive,
+        loc: GeneratedSource,
+      },
+      property: dep.path[idx].property,
+      loc: GeneratedSource,
+    });
+    lowerValueToTemporary(builder, {
+      kind: 'StoreLocal',
+      lvalue: {kind: InstructionKind.Const, place: {...returnPlace}},
+      value: {...tmpConsequent},
+      type: null,
+      loc: GeneratedSource,
+    });
+    return {
+      kind: 'goto',
+      variant: GotoVariant.Break,
+      block: continuationBlock.id,
+      id: makeInstructionId(0),
+      loc: GeneratedSource,
+    };
+  });
+  builder.terminateWithContinuation(
+    {
+      kind: 'optional',
+      optional: dep.path[idx].optional,
+      test: testBlock,
+      fallthrough: continuationBlock.id,
+      id: makeInstructionId(0),
+      loc: GeneratedSource,
+    },
+    continuationBlock,
+  );
+
+  return returnPlace.identifier;
+}

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

@@ -10,7 +10,6 @@ import {CompilerError, SourceLocation} from '..';
 import {
   ArrayExpression,
   Effect,
-  Environment,
   FunctionExpression,
   GeneratedSource,
   HIRFunction,
@@ -29,6 +28,9 @@ import {
   isSetStateType,
   isFireFunctionType,
   makeScopeId,
+  HIR,
+  BasicBlock,
+  BlockId,
 } from '../HIR';
 import {collectHoistablePropertyLoadsInInnerFn} from '../HIR/CollectHoistablePropertyLoads';
 import {collectOptionalChainSidemap} from '../HIR/CollectOptionalChainDependencies';
@@ -44,7 +46,12 @@ import {
   DependencyCollectionContext,
   handleInstruction,
 } from '../HIR/PropagateScopeDependenciesHIR';
-import {eachInstructionOperand, eachTerminalOperand} from '../HIR/visitors';
+import {buildDependencyInstructions} from '../HIR/ScopeDependencyUtils';
+import {
+  eachInstructionOperand,
+  eachTerminalOperand,
+  terminalFallthrough,
+} from '../HIR/visitors';
 import {empty} from '../Utils/Stack';
 import {getOrInsertWith} from '../Utils/utils';
 
@@ -53,7 +60,6 @@ import {getOrInsertWith} from '../Utils/utils';
  * a second argument to the useEffect call if no dependency array is provided.
  */
 export function inferEffectDependencies(fn: HIRFunction): void {
-  let hasRewrite = false;
   const fnExpressions = new Map<
     IdentifierId,
     TInstruction<FunctionExpression>
@@ -86,6 +92,7 @@ export function inferEffectDependencies(fn: HIRFunction): void {
    * reactive(Identifier i) = Union_{reference of i}(reactive(reference))
    */
   const reactiveIds = inferReactiveIdentifiers(fn);
+  const rewriteBlocks: Array<BasicBlock> = [];
 
   for (const [, block] of fn.body.blocks) {
     if (block.terminal.kind === 'scope') {
@@ -101,7 +108,7 @@ export function inferEffectDependencies(fn: HIRFunction): void {
         );
       }
     }
-    const rewriteInstrs = new Map<InstructionId, Array<Instruction>>();
+    const rewriteInstrs: Array<SpliceInfo> = [];
     for (const instr of block.instructions) {
       const {value, lvalue} = instr;
       if (value.kind === 'FunctionExpression') {
@@ -165,7 +172,6 @@ export function inferEffectDependencies(fn: HIRFunction): void {
         ) {
           // We have a useEffect call with no deps array, so we need to infer the deps
           const effectDeps: Array<Place> = [];
-          const newInstructions: Array<Instruction> = [];
           const deps: ArrayExpression = {
             kind: 'ArrayExpression',
             elements: effectDeps,
@@ -196,24 +202,28 @@ export function inferEffectDependencies(fn: HIRFunction): void {
              */
 
             const usedDeps = [];
-            for (const dep of minimalDeps) {
+            for (const maybeDep of minimalDeps) {
               if (
-                ((isUseRefType(dep.identifier) ||
-                  isSetStateType(dep.identifier)) &&
-                  !reactiveIds.has(dep.identifier.id)) ||
-                isFireFunctionType(dep.identifier)
+                ((isUseRefType(maybeDep.identifier) ||
+                  isSetStateType(maybeDep.identifier)) &&
+                  !reactiveIds.has(maybeDep.identifier.id)) ||
+                isFireFunctionType(maybeDep.identifier)
               ) {
                 // exclude non-reactive hook results, which will never be in a memo block
                 continue;
               }
 
-              const {place, instructions} = writeDependencyToInstructions(
+              const dep = truncateDepAtCurrent(maybeDep);
+              const {place, value, exitBlockId} = buildDependencyInstructions(
                 dep,
-                reactiveIds.has(dep.identifier.id),
                 fn.env,
-                fnExpr.loc,
               );
-              newInstructions.push(...instructions);
+              rewriteInstrs.push({
+                kind: 'block',
+                location: instr.id,
+                value,
+                exitBlockId: exitBlockId,
+              });
               effectDeps.push(place);
               usedDeps.push(dep);
             }
@@ -234,27 +244,32 @@ export function inferEffectDependencies(fn: HIRFunction): void {
               });
             }
 
-            newInstructions.push({
-              id: makeInstructionId(0),
-              loc: GeneratedSource,
-              lvalue: {...depsPlace, effect: Effect.Mutate},
-              value: deps,
-            });
-
             // Step 2: push the inferred deps array as an argument of the useEffect
+            rewriteInstrs.push({
+              kind: 'instr',
+              location: instr.id,
+              value: {
+                id: makeInstructionId(0),
+                loc: GeneratedSource,
+                lvalue: {...depsPlace, effect: Effect.Mutate},
+                value: deps,
+              },
+            });
             value.args.push({...depsPlace, effect: Effect.Freeze});
-            rewriteInstrs.set(instr.id, newInstructions);
             fn.env.inferredEffectLocations.add(callee.loc);
           } else if (loadGlobals.has(value.args[0].identifier.id)) {
             // Global functions have no reactive dependencies, so we can insert an empty array
-            newInstructions.push({
-              id: makeInstructionId(0),
-              loc: GeneratedSource,
-              lvalue: {...depsPlace, effect: Effect.Mutate},
-              value: deps,
+            rewriteInstrs.push({
+              kind: 'instr',
+              location: instr.id,
+              value: {
+                id: makeInstructionId(0),
+                loc: GeneratedSource,
+                lvalue: {...depsPlace, effect: Effect.Mutate},
+                value: deps,
+              },
             });
             value.args.push({...depsPlace, effect: Effect.Freeze});
-            rewriteInstrs.set(instr.id, newInstructions);
             fn.env.inferredEffectLocations.add(callee.loc);
           }
         } else if (
@@ -285,21 +300,13 @@ export function inferEffectDependencies(fn: HIRFunction): void {
         }
       }
     }
-    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;
-    }
+    rewriteSplices(block, rewriteInstrs, rewriteBlocks);
   }
-  if (hasRewrite) {
+
+  if (rewriteBlocks.length > 0) {
+    for (const block of rewriteBlocks) {
+      fn.body.blocks.set(block.id, block);
+    }
     // Renumber instructions and fix scope ranges
     markInstructionIds(fn.body);
     fixScopeAndIdentifierRanges(fn.body);
@@ -307,63 +314,144 @@ export function inferEffectDependencies(fn: HIRFunction): void {
   }
 }
 
-function writeDependencyToInstructions(
+function truncateDepAtCurrent(
   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;
+): ReactiveScopeDependency {
+  const idx = dep.path.findIndex(path => path.property === 'current');
+  if (idx === -1) {
+    return dep;
+  } else {
+    return {...dep, path: dep.path.slice(0, idx)};
+  }
+}
+
+type SpliceInfo =
+  | {kind: 'instr'; location: InstructionId; value: Instruction}
+  | {
+      kind: 'block';
+      location: InstructionId;
+      value: HIR;
+      exitBlockId: BlockId;
+    };
+
+function rewriteSplices(
+  originalBlock: BasicBlock,
+  splices: Array<SpliceInfo>,
+  rewriteBlocks: Array<BasicBlock>,
+): void {
+  if (splices.length === 0) {
+    return;
+  }
+  /**
+   * Splice instructions or value blocks into the original block.
+   * --- original block ---
+   * bb_original
+   *   instr1
+   *   ...
+   *   instr2 <-- splice location
+   *   instr3
+   *   ...
+   *   <original terminal>
+   *
+   * If there is more than one block in the splice, this means that we're
+   * splicing in a set of value-blocks of the following structure:
+   * --- blocks we're splicing in ---
+   * bb_entry:
+   *   instrEntry
+   *   ...
+   *   <splice terminal> fallthrough=bb_exit
+   *
+   * bb1(value):
+   *   ...
+   *
+   * bb_exit:
+   *   instrExit
+   *   ...
+   *   <synthetic terminal>
+   *
+   *
+   * --- rewritten blocks ---
+   * bb_original
+   *   instr1
+   *   ... (original instructions)
+   *   instr2
+   *   instrEntry
+   *   ... (spliced instructions)
+   *   <splice terminal> fallthrough=bb_exit
+   *
+   * bb1(value):
+   *   ...
+   *
+   * bb_exit:
+   *   instrExit
+   *   ... (spliced instructions)
+   *   instr3
+   *   ... (original instructions)
+   *   <original terminal>
+   */
+  const originalInstrs = originalBlock.instructions;
+  let currBlock: BasicBlock = {...originalBlock, instructions: []};
+  rewriteBlocks.push(currBlock);
+
+  let cursor = 0;
+  for (const rewrite of splices) {
+    while (originalInstrs[cursor].id < rewrite.location) {
+      CompilerError.invariant(
+        originalInstrs[cursor].id < originalInstrs[cursor + 1].id,
+        {
+          reason:
+            '[InferEffectDependencies] Internal invariant broken: expected block instructions to be sorted',
+          loc: originalInstrs[cursor].loc,
+        },
+      );
+      currBlock.instructions.push(originalInstrs[cursor]);
+      cursor++;
     }
-    if (path.property === 'current') {
-      /*
-       * Prune ref.current accesses. This may over-capture for non-ref values with
-       * a current property, but that's fine.
-       */
-      break;
+    CompilerError.invariant(originalInstrs[cursor].id === rewrite.location, {
+      reason:
+        '[InferEffectDependencies] Internal invariant broken: splice location not found',
+      loc: originalInstrs[cursor].loc,
+    });
+
+    if (rewrite.kind === 'instr') {
+      currBlock.instructions.push(rewrite.value);
+    } else {
+      const {entry, blocks} = rewrite.value;
+      const entryBlock = blocks.get(entry)!;
+      // splice in all instructions from the entry block
+      currBlock.instructions.push(...entryBlock.instructions);
+      if (blocks.size > 1) {
+        /**
+         * We're splicing in a set of value-blocks, which means we need
+         * to push new blocks and update terminals.
+         */
+        CompilerError.invariant(
+          terminalFallthrough(entryBlock.terminal) === rewrite.exitBlockId,
+          {
+            reason:
+              '[InferEffectDependencies] Internal invariant broken: expected entry block to have a fallthrough',
+            loc: entryBlock.terminal.loc,
+          },
+        );
+        const originalTerminal = currBlock.terminal;
+        currBlock.terminal = entryBlock.terminal;
+
+        for (const [id, block] of blocks) {
+          if (id === entry) {
+            continue;
+          }
+          if (id === rewrite.exitBlockId) {
+            block.terminal = originalTerminal;
+            currBlock = block;
+          }
+          rewriteBlocks.push(block);
+        }
+      }
     }
-    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;
+  }
+  for (let i = cursor; i < originalInstrs.length; i++) {
+    currBlock.instructions.push(originalInstrs[i]);
   }
-  currValue.effect = Effect.Freeze;
-  return {place: currValue, instructions};
 }
 
 function inferReactiveIdentifiers(fn: HIRFunction): Set<IdentifierId> {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/reactive-optional-chain.expect.md

@@ -8,10 +8,16 @@ import {print} from 'shared-runtime';
 
 // TODO: take optional chains as dependencies
 function ReactiveMemberExpr({cond, propVal}) {
-  const obj = {a: cond ? {b: propVal} : null};
+  const obj = {a: cond ? {b: propVal} : null, c: null};
   useEffect(() => print(obj.a?.b));
+  useEffect(() => print(obj.c?.d));
 }
 
+export const FIXTURE_ENTRYPOINT = {
+  fn: ReactiveMemberExpr,
+  params: [{cond: true, propVal: 1}],
+};
+
 ```
 
 ## Code
@@ -23,7 +29,7 @@ import { print } from "shared-runtime";
 
 // TODO: take optional chains as dependencies
 function ReactiveMemberExpr(t0) {
-  const $ = _c(7);
+  const $ = _c(9);
   const { cond, propVal } = t0;
   let t1;
   if ($[0] !== cond || $[1] !== propVal) {
@@ -36,7 +42,7 @@ function ReactiveMemberExpr(t0) {
   }
   let t2;
   if ($[3] !== t1) {
-    t2 = { a: t1 };
+    t2 = { a: t1, c: null };
     $[3] = t1;
     $[4] = t2;
   } else {
@@ -51,10 +57,25 @@ function ReactiveMemberExpr(t0) {
   } else {
     t3 = $[6];
   }
-  useEffect(t3, [obj.a]);
+  useEffect(t3, [obj.a?.b]);
+  let t4;
+  if ($[7] !== obj.c?.d) {
+    t4 = () => print(obj.c?.d);
+    $[7] = obj.c?.d;
+    $[8] = t4;
+  } else {
+    t4 = $[8];
+  }
+  useEffect(t4, [obj.c?.d]);
 }
 
+export const FIXTURE_ENTRYPOINT = {
+  fn: ReactiveMemberExpr,
+  params: [{ cond: true, propVal: 1 }],
+};
+
 ```
       
 ### Eval output
-(kind: exception) Fixture not implemented
+(kind: ok) 
+logs: [1,undefined]

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/reactive-optional-chain.js

@@ -4,6 +4,12 @@ import {print} from 'shared-runtime';
 
 // TODO: take optional chains as dependencies
 function ReactiveMemberExpr({cond, propVal}) {
-  const obj = {a: cond ? {b: propVal} : null};
+  const obj = {a: cond ? {b: propVal} : null, c: null};
   useEffect(() => print(obj.a?.b));
+  useEffect(() => print(obj.c?.d));
 }
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: ReactiveMemberExpr,
+  params: [{cond: true, propVal: 1}],
+};