[compiler][hir-rewrite] Stop hoisting PropertyLoads out of function d…

…ecls

ghstack-source-id: da8ae28a969cd027c076e36a9a815665a3af5754
Pull Request resolved: #31066

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

@@ -11,6 +11,7 @@ import {
   BasicBlock,
   BlockId,
   DependencyPathEntry,
+  FunctionExpression,
   GeneratedSource,
   HIRFunction,
   Identifier,
@@ -19,6 +20,11 @@ import {
   ReactiveScopeDependency,
   ScopeId,
 } from './HIR';
+import {
+  eachInstructionOperand,
+  eachInstructionValueLValue,
+  eachInstructionValueOperand,
+} from './visitors';
 
 /**
  * Helper function for `PropagateScopeDependencies`.
@@ -76,7 +82,12 @@ export function collectHoistablePropertyLoads(
 ): ReadonlyMap<ScopeId, BlockInfo> {
   const tree = new Tree();
 
-  const nodes = collectNonNullsInBlocks(fn, temporaries, optionals, tree);
+  const functionExpressionReferences = collectFunctionExpressionRValues(fn);
+  const realTemporaries = new Map(temporaries);
+  for (const reference of functionExpressionReferences) {
+    realTemporaries.delete(reference);
+  }
+  const nodes = collectNonNullsInBlocks(fn, realTemporaries, optionals, tree);
   propagateNonNull(fn, nodes, tree);
 
   const nodesKeyedByScopeId = new Map<ScopeId, BlockInfo>();
@@ -506,3 +517,38 @@ function reduceMaybeOptionalChains(
     }
   } while (changed);
 }
+
+function collectFunctionExpressionRValues(fn: HIRFunction): Set<IdentifierId> {
+  const uses = new Map<IdentifierId, Set<IdentifierId>>();
+  const functionExpressionReferences = new Set<IdentifierId>();
+
+  for (const [_, block] of fn.body.blocks) {
+    for (const {lvalue, value} of block.instructions) {
+      if (value.kind === 'FunctionExpression') {
+        for (const reference of value.loweredFunc.dependencies) {
+          const queued = [reference.identifier.id];
+          while (queued.length > 0) {
+            const top = queued.pop()!;
+            functionExpressionReferences.add(top);
+            for (const reference of uses.get(top) ?? []) {
+              queued.push(reference);
+            }
+          }
+        }
+      } else {
+        const operands = [
+          ...eachInstructionValueOperand(value),
+          ...eachInstructionValueLValue(value),
+        ];
+        if (!operands.every(operand => operand.identifier.name == null)) {
+          continue;
+        }
+        uses.set(
+          lvalue.identifier.id,
+          new Set(operands.map(operand => operand.identifier.id)),
+        );
+      }
+    }
+  }
+  return functionExpressionReferences;
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/function-expression-dependencies-not-hoistable.expect.md

@@ -0,0 +1,75 @@
+
+## Input
+
+```javascript
+// @enablePropagateDepsInHIR
+
+import {makeArray, identity, useIdentity} from 'shared-runtime';
+
+function useFoo(a) {
+  const fn = () => a.b.c;
+  useIdentity();
+  const x = makeArray();
+  x.push(identity(a?.b.c));
+  return [fn, x];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [null],
+  sequentialRenders: [null, {b: {c: 4}}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enablePropagateDepsInHIR
+
+import { makeArray, identity, useIdentity } from "shared-runtime";
+
+function useFoo(a) {
+  const $ = _c(7);
+  let t0;
+  if ($[0] !== a) {
+    t0 = () => a.b.c;
+    $[0] = a;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  const fn = t0;
+  useIdentity();
+  let x;
+  if ($[2] !== a?.b.c) {
+    x = makeArray();
+    x.push(identity(a?.b.c));
+    $[2] = a?.b.c;
+    $[3] = x;
+  } else {
+    x = $[3];
+  }
+  let t1;
+  if ($[4] !== fn || $[5] !== x) {
+    t1 = [fn, x];
+    $[4] = fn;
+    $[5] = x;
+    $[6] = t1;
+  } else {
+    t1 = $[6];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [null],
+  sequentialRenders: [null, { b: { c: 4 } }],
+};
+
+```
+      
+### Eval output
+(kind: ok) ["[[ function params=0 ]]",[null]]
+["[[ function params=0 ]]",[4]]

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/function-expression-dependencies-not-hoistable.ts

@@ -0,0 +1,17 @@
+// @enablePropagateDepsInHIR
+
+import {makeArray, identity, useIdentity} from 'shared-runtime';
+
+function useFoo(a) {
+  const fn = () => a.b.c;
+  useIdentity();
+  const x = makeArray();
+  x.push(identity(a?.b.c));
+  return [fn, x];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [null],
+  sequentialRenders: [null, {b: {c: 4}}],
+};