[compiler] Context variables as dependencies (facebook#31582)

We previously didn't track context variables in the hoistable values
sidemap of `propagateScopeDependencies`. This was overly conservative as
we *do* track the mutable range of context variables, and it is safe to
hoist accesses to context variables after their last direct / aliased
maybe-assignment.

```js
function Component({value}) {
  // start of mutable range for `x`
  let x = DEFAULT;
  const setX = () => x = value;
  const aliasedSet = maybeAlias(setX);
  maybeCall(aliasedSet);
  // end of mutable range for `x`

  // here, we should be able to take x (and property reads
  // off of x) as dependencies
  return <Jsx value={x} />
}
```
---
[//]: # (BEGIN SAPLING FOOTER)
Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with [ReviewStack](https://reviewstack.dev/facebook/react/pull/31582).
* facebook#31583
* __->__ facebook#31582

Author: mofeiZ

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

@@ -840,6 +840,11 @@ export type LoadLocal = {
   place: Place;
   loc: SourceLocation;
 };
+export type LoadContext = {
+  kind: 'LoadContext';
+  place: Place;
+  loc: SourceLocation;
+};
 
 /*
  * The value of a given instruction. Note that values are not recursive: complex
@@ -852,11 +857,7 @@ export type LoadLocal = {
 
 export type InstructionValue =
   | LoadLocal
-  | {
-      kind: 'LoadContext';
-      place: Place;
-      loc: SourceLocation;
-    }
+  | LoadContext
   | {
       kind: 'DeclareLocal';
       lvalue: LValue;

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

@@ -17,6 +17,11 @@ import {
   areEqualPaths,
   IdentifierId,
   Terminal,
+  InstructionValue,
+  LoadContext,
+  TInstruction,
+  FunctionExpression,
+  ObjectMethod,
 } from './HIR';
 import {
   collectHoistablePropertyLoads,
@@ -223,11 +228,25 @@ export function collectTemporariesSidemap(
     fn,
     usedOutsideDeclaringScope,
     temporaries,
-    false,
+    null,
   );
   return temporaries;
 }
 
+function isLoadContextMutable(
+  instrValue: InstructionValue,
+  id: InstructionId,
+): instrValue is LoadContext {
+  if (instrValue.kind === 'LoadContext') {
+    CompilerError.invariant(instrValue.place.identifier.scope != null, {
+      reason:
+        '[PropagateScopeDependencies] Expected all context variables to be assigned a scope',
+      loc: instrValue.loc,
+    });
+    return id >= instrValue.place.identifier.scope.range.end;
+  }
+  return false;
+}
 /**
  * Recursive collect a sidemap of all `LoadLocal` and `PropertyLoads` with a
  * function and all nested functions.
@@ -239,17 +258,21 @@ function collectTemporariesSidemapImpl(
   fn: HIRFunction,
   usedOutsideDeclaringScope: ReadonlySet<DeclarationId>,
   temporaries: Map<IdentifierId, ReactiveScopeDependency>,
-  isInnerFn: boolean,
+  innerFnContext: {instrId: InstructionId} | null,
 ): void {
   for (const [_, block] of fn.body.blocks) {
-    for (const instr of block.instructions) {
-      const {value, lvalue} = instr;
+    for (const {value, lvalue, id: origInstrId} of block.instructions) {
+      const instrId =
+        innerFnContext != null ? innerFnContext.instrId : origInstrId;
       const usedOutside = usedOutsideDeclaringScope.has(
         lvalue.identifier.declarationId,
       );
 
       if (value.kind === 'PropertyLoad' && !usedOutside) {
-        if (!isInnerFn || temporaries.has(value.object.identifier.id)) {
+        if (
+          innerFnContext == null ||
+          temporaries.has(value.object.identifier.id)
+        ) {
           /**
            * All dependencies of a inner / nested function must have a base
            * identifier from the outermost component / hook. This is because the
@@ -265,13 +288,13 @@ function collectTemporariesSidemapImpl(
           temporaries.set(lvalue.identifier.id, property);
         }
       } else if (
-        value.kind === 'LoadLocal' &&
+        (value.kind === 'LoadLocal' || isLoadContextMutable(value, instrId)) &&
         lvalue.identifier.name == null &&
         value.place.identifier.name !== null &&
         !usedOutside
       ) {
         if (
-          !isInnerFn ||
+          innerFnContext == null ||
           fn.context.some(
             context => context.identifier.id === value.place.identifier.id,
           )
@@ -289,7 +312,7 @@ function collectTemporariesSidemapImpl(
           value.loweredFunc.func,
           usedOutsideDeclaringScope,
           temporaries,
-          true,
+          innerFnContext ?? {instrId},
         );
       }
     }
@@ -364,7 +387,7 @@ class Context {
    * Tracks the traversal state. See Context.declare for explanation of why this
    * is needed.
    */
-  inInnerFn: boolean = false;
+  #innerFnContext: {outerInstrId: InstructionId} | null = null;
 
   constructor(
     temporariesUsedOutsideScope: ReadonlySet<DeclarationId>,
@@ -434,7 +457,7 @@ class Context {
    *     by root identifier mutable ranges).
    */
   declare(identifier: Identifier, decl: Decl): void {
-    if (this.inInnerFn) return;
+    if (this.#innerFnContext != null) return;
     if (!this.#declarations.has(identifier.declarationId)) {
       this.#declarations.set(identifier.declarationId, decl);
     }
@@ -577,11 +600,14 @@ class Context {
       currentScope.reassignments.add(place.identifier);
     }
   }
-  enterInnerFn<T>(cb: () => T): T {
-    const wasInInnerFn = this.inInnerFn;
-    this.inInnerFn = true;
+  enterInnerFn<T>(
+    innerFn: TInstruction<FunctionExpression> | TInstruction<ObjectMethod>,
+    cb: () => T,
+  ): T {
+    const prevContext = this.#innerFnContext;
+    this.#innerFnContext = this.#innerFnContext ?? {outerInstrId: innerFn.id};
     const result = cb();
-    this.inInnerFn = wasInInnerFn;
+    this.#innerFnContext = prevContext;
     return result;
   }
 
@@ -724,9 +750,14 @@ function collectDependencies(
            * Recursively visit the inner function to extract dependencies there
            */
           const innerFn = instr.value.loweredFunc.func;
-          context.enterInnerFn(() => {
-            handleFunction(innerFn);
-          });
+          context.enterInnerFn(
+            instr as
+              | TInstruction<FunctionExpression>
+              | TInstruction<ObjectMethod>,
+            () => {
+              handleFunction(innerFn);
+            },
+          );
         } else {
           handleInstruction(instr, context);
         }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-functiondecl-hoisting.expect.md

@@ -58,18 +58,16 @@ function Foo(t0) {
     bar = $[1];
     result = $[2];
   }
-
-  const t1 = bar;
-  let t2;
-  if ($[3] !== result || $[4] !== t1) {
-    t2 = <Stringify result={result} fn={t1} shouldInvokeFns={true} />;
-    $[3] = result;
-    $[4] = t1;
-    $[5] = t2;
+  let t1;
+  if ($[3] !== bar || $[4] !== result) {
+    t1 = <Stringify result={result} fn={bar} shouldInvokeFns={true} />;
+    $[3] = bar;
+    $[4] = result;
+    $[5] = t1;
   } else {
-    t2 = $[5];
+    t1 = $[5];
   }
-  return t2;
+  return t1;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/destructure-array-assignment-to-context-var.expect.md

@@ -43,16 +43,15 @@ function Component(props) {
   } else {
     x = $[1];
   }
-  const t0 = x;
-  let t1;
-  if ($[2] !== t0) {
-    t1 = { x: t0 };
-    $[2] = t0;
-    $[3] = t1;
+  let t0;
+  if ($[2] !== x) {
+    t0 = { x };
+    $[2] = x;
+    $[3] = t0;
   } else {
-    t1 = $[3];
+    t0 = $[3];
   }
-  return t1;
+  return t0;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/destructure-array-declaration-to-context-var.expect.md

@@ -42,16 +42,15 @@ function Component(props) {
   } else {
     x = $[1];
   }
-  const t0 = x;
-  let t1;
-  if ($[2] !== t0) {
-    t1 = <div>{t0}</div>;
-    $[2] = t0;
-    $[3] = t1;
+  let t0;
+  if ($[2] !== x) {
+    t0 = <div>{x}</div>;
+    $[2] = x;
+    $[3] = t0;
   } else {
-    t1 = $[3];
+    t0 = $[3];
   }
-  return t1;
+  return t0;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/destructure-object-assignment-to-context-var.expect.md

@@ -43,16 +43,15 @@ function Component(props) {
   } else {
     x = $[1];
   }
-  const t0 = x;
-  let t1;
-  if ($[2] !== t0) {
-    t1 = { x: t0 };
-    $[2] = t0;
-    $[3] = t1;
+  let t0;
+  if ($[2] !== x) {
+    t0 = { x };
+    $[2] = x;
+    $[3] = t0;
   } else {
-    t1 = $[3];
+    t0 = $[3];
   }
-  return t1;
+  return t0;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/destructure-object-declaration-to-context-var.expect.md

@@ -42,16 +42,15 @@ function Component(props) {
   } else {
     x = $[1];
   }
-  const t0 = x;
-  let t1;
-  if ($[2] !== t0) {
-    t1 = { x: t0 };
-    $[2] = t0;
-    $[3] = t1;
+  let t0;
+  if ($[2] !== x) {
+    t0 = { x };
+    $[2] = x;
+    $[3] = t0;
   } else {
-    t1 = $[3];
+    t0 = $[3];
   }
-  return t1;
+  return t0;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lambda-mutated-non-reactive-to-reactive.expect.md

@@ -33,17 +33,15 @@ function f(a) {
   } else {
     x = $[1];
   }
-
-  const t0 = x;
-  let t1;
-  if ($[2] !== t0) {
-    t1 = <div x={t0} />;
-    $[2] = t0;
-    $[3] = t1;
+  let t0;
+  if ($[2] !== x) {
+    t0 = <div x={x} />;
+    $[2] = x;
+    $[3] = t0;
   } else {
-    t1 = $[3];
+    t0 = $[3];
   }
-  return t1;
+  return t0;
 }
 
 export const FIXTURE_ENTRYPOINT = {