[compiler] bugfix for hoistable deps for nested functions

`PropertyPathRegistry` is responsible for uniqueing identifier and property paths. This is necessary for the hoistability CFG merging logic which takes unions and intersections of these nodes to determine a basic block's hoistable reads, as a function of its neighbors. We also depend on this to merge optional chained and non-optional chained property paths

This fixes a small bug in #31066 in which we create a new registry for nested functions. Now, we use the same registry for a component / hook and all its inner functions

'

Author: mofeiZ

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

@@ -1,5 +1,6 @@
 import {CompilerError} from '../CompilerError';
 import {inRange} from '../ReactiveScopes/InferReactiveScopeVariables';
+import {printDependency} from '../ReactiveScopes/PrintReactiveFunction';
 import {
   Set_equal,
   Set_filter,
@@ -23,6 +24,8 @@ import {
 } from './HIR';
 import {collectTemporariesSidemap} from './PropagateScopeDependenciesHIR';
 
+const DEBUG_PRINT = false;
+
 /**
  * Helper function for `PropagateScopeDependencies`. Uses control flow graph
  * analysis to determine which `Identifier`s can be assumed to be non-null
@@ -86,15 +89,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 +107,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 +276,7 @@ class PropertyPathRegistry {
 
 function getMaybeNonNullInInstruction(
   instr: InstructionValue,
-  context: CollectNonNullsInBlocksContext,
+  context: CollectHoistablePropertyLoadsContext,
 ): PropertyPathNode | null {
   let path = null;
   if (instr.kind === 'PropertyLoad') {
@@ -262,7 +295,7 @@ function getMaybeNonNullInInstruction(
 function isImmutableAtInstr(
   identifier: Identifier,
   instr: InstructionId,
-  context: CollectNonNullsInBlocksContext,
+  context: CollectHoistablePropertyLoadsContext,
 ): boolean {
   if (context.nestedFnImmutableContext != null) {
     /**
@@ -295,22 +328,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
@@ -358,18 +378,22 @@ function collectNonNullsInBlocks(
           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,
+            temporaries: innerTemporaries, // TODO: remove in later PR
+            hoistableFromOptionals: innerOptionals.hoistableObjects, // TODO: remove in later PR
+            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),

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

@@ -46,7 +46,7 @@ export function propagateScopeDependenciesHIR(fn: HIRFunction): void {
 
   const hoistablePropertyLoads = keyByScopeId(
     fn,
-    collectHoistablePropertyLoads(fn, temporaries, hoistableObjects, null),
+    collectHoistablePropertyLoads(fn, temporaries, hoistableObjects),
   );
 
   const scopeDeps = collectDependencies(

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/repro-invariant.expect.md

@@ -0,0 +1,61 @@
+
+## Input
+
+```javascript
+// @enablePropagateDepsInHIR
+import {Stringify} from 'shared-runtime';
+
+function Foo({data}) {
+  return (
+    <Stringify foo={() => data.a.d} bar={data.a?.b.c} shouldInvokeFns={true} />
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{data: {a: null}}],
+  sequentialRenders: [{data: {a: {b: {c: 4}}}}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enablePropagateDepsInHIR
+import { Stringify } from "shared-runtime";
+
+function Foo(t0) {
+  const $ = _c(5);
+  const { data } = t0;
+  let t1;
+  if ($[0] !== data.a.d) {
+    t1 = () => data.a.d;
+    $[0] = data.a.d;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  const t2 = data.a?.b.c;
+  let t3;
+  if ($[2] !== t1 || $[3] !== t2) {
+    t3 = <Stringify foo={t1} bar={t2} shouldInvokeFns={true} />;
+    $[2] = t1;
+    $[3] = t2;
+    $[4] = t3;
+  } else {
+    t3 = $[4];
+  }
+  return t3;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{ data: { a: null } }],
+  sequentialRenders: [{ data: { a: { b: { c: 4 } } } }],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"foo":{"kind":"Function"},"bar":4,"shouldInvokeFns":true}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/repro-invariant.tsx

@@ -0,0 +1,14 @@
+// @enablePropagateDepsInHIR
+import {Stringify} from 'shared-runtime';
+
+function Foo({data}) {
+  return (
+    <Stringify foo={() => data.a.d} bar={data.a?.b.c} shouldInvokeFns={true} />
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{data: {a: null}}],
+  sequentialRenders: [{data: {a: {b: {c: 4}}}}],
+};