[compiler][hir] Only hoist always-accessed PropertyLoads from function decls

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

Author: mofeiZ

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

@@ -7,6 +7,7 @@ import {
   Set_union,
   getOrInsertDefault,
 } from '../Utils/utils';
+import {collectOptionalChainSidemap} from './CollectOptionalChainDependencies';
 import {
   BasicBlock,
   BlockId,
@@ -19,6 +20,7 @@ import {
   ReactiveScopeDependency,
   ScopeId,
 } from './HIR';
+import {collectTemporariesSidemap} from './PropagateScopeDependenciesHIR';
 
 /**
  * Helper function for `PropagateScopeDependencies`.
@@ -73,23 +75,38 @@ export function collectHoistablePropertyLoads(
   fn: HIRFunction,
   temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
   optionals: ReadonlyMap<BlockId, ReactiveScopeDependency>,
-): ReadonlyMap<ScopeId, BlockInfo> {
+): ReadonlyMap<BlockId, BlockInfo> {
   const registry = new PropertyPathRegistry();
 
-  const nodes = collectNonNullsInBlocks(fn, temporaries, optionals, registry);
+  const functionExpressionReferences = collectFunctionExpressionRValues(fn);
+  const reallyAccessedTemporaries = new Map(
+    [...temporaries].filter(([id]) => !functionExpressionReferences.has(id)),
+  );
+  const nodes = collectNonNullsInBlocks(
+    fn,
+    reallyAccessedTemporaries,
+    optionals,
+    registry,
+  );
   propagateNonNull(fn, nodes, registry);
 
-  const nodesKeyedByScopeId = new Map<ScopeId, BlockInfo>();
+  return nodes;
+}
+
+export function keyByScopeId<T>(
+  fn: HIRFunction,
+  source: ReadonlyMap<BlockId, T>,
+): ReadonlyMap<ScopeId, T> {
+  const keyedByScopeId = new Map<ScopeId, T>();
   for (const [_, block] of fn.body.blocks) {
     if (block.terminal.kind === 'scope') {
-      nodesKeyedByScopeId.set(
+      keyedByScopeId.set(
         block.terminal.scope.id,
-        nodes.get(block.terminal.block)!,
+        source.get(block.terminal.block)!,
       );
     }
   }
-
-  return nodesKeyedByScopeId;
+  return keyedByScopeId;
 }
 
 export type BlockInfo = {
@@ -317,6 +334,27 @@ function collectNonNullsInBlocks(
             assumedNonNullObjects,
           );
         }
+      } else if (
+        instr.value.kind === 'FunctionExpression' &&
+        !fn.env.config.enableTreatFunctionDepsAsConditional
+      ) {
+        const innerFn = instr.value.loweredFunc;
+        const innerTemporaries = collectTemporariesSidemap(
+          innerFn.func,
+          new Set(),
+        );
+        const optionals = collectOptionalChainSidemap(innerFn.func);
+        const innerHoistableMap = collectHoistablePropertyLoads(
+          innerFn.func,
+          innerTemporaries,
+          optionals.hoistableObjects,
+        );
+        const innerHoistables = assertNonNull(
+          innerHoistableMap.get(innerFn.func.body.entry),
+        );
+        for (const entry of innerHoistables.assumedNonNullObjects) {
+          assumedNonNullObjects.add(entry);
+        }
       }
     }
   }
@@ -518,3 +556,25 @@ function reduceMaybeOptionalChains(
     }
   } while (changed);
 }
+
+function collectFunctionExpressionRValues(fn: HIRFunction): Set<IdentifierId> {
+  const sources = new Map<IdentifierId, 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) {
+          let curr: IdentifierId | undefined = reference.identifier.id;
+          while (curr != null) {
+            functionExpressionReferences.add(curr);
+            curr = sources.get(curr);
+          }
+        }
+      } else if (value.kind === 'PropertyLoad') {
+        sources.set(lvalue.identifier.id, value.object.identifier.id);
+      }
+    }
+  }
+  return functionExpressionReferences;
+}

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

@@ -17,7 +17,10 @@ import {
   areEqualPaths,
   IdentifierId,
 } from './HIR';
-import {collectHoistablePropertyLoads} from './CollectHoistablePropertyLoads';
+import {
+  collectHoistablePropertyLoads,
+  keyByScopeId,
+} from './CollectHoistablePropertyLoads';
 import {
   ScopeBlockTraversal,
   eachInstructionOperand,
@@ -41,10 +44,9 @@ export function propagateScopeDependenciesHIR(fn: HIRFunction): void {
     hoistableObjects,
   } = collectOptionalChainSidemap(fn);
 
-  const hoistablePropertyLoads = collectHoistablePropertyLoads(
+  const hoistablePropertyLoads = keyByScopeId(
     fn,
-    temporaries,
-    hoistableObjects,
+    collectHoistablePropertyLoads(fn, temporaries, hoistableObjects),
   );
 
   const scopeDeps = collectDependencies(
@@ -209,7 +211,7 @@ function findTemporariesUsedOutsideDeclaringScope(
  * of $1, as the evaluation of `arr.length` changes between instructions $1 and
  * $3. We do not track $1 -> arr.length in this case.
  */
-function collectTemporariesSidemap(
+export function collectTemporariesSidemap(
   fn: HIRFunction,
   usedOutsideDeclaringScope: ReadonlySet<DeclarationId>,
 ): ReadonlyMap<IdentifierId, ReactiveScopeDependency> {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/reduce-reactive-deps/infer-function-cond-access-not-hoisted.expect.md

@@ -0,0 +1,75 @@
+
+## Input
+
+```javascript
+// @enablePropagateDepsInHIR
+
+import { Stringify } from "shared-runtime";
+
+function Foo({a, shouldReadA}) {
+  return <Stringify fn={() => {
+    if (shouldReadA) return a.b.c;
+    return null;
+  }} shouldInvokeFns={true} />
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{a: null, shouldReadA: true}],
+  sequentialRenders: [
+    {a: null, shouldReadA: true},
+    {a: null, shouldReadA: false},
+    {a: {b: {c: 4}}, shouldReadA: true},
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enablePropagateDepsInHIR
+
+import { Stringify } from "shared-runtime";
+
+function Foo(t0) {
+  const $ = _c(3);
+  const { a, shouldReadA } = t0;
+  let t1;
+  if ($[0] !== shouldReadA || $[1] !== a) {
+    t1 = (
+      <Stringify
+        fn={() => {
+          if (shouldReadA) {
+            return a.b.c;
+          }
+          return null;
+        }}
+        shouldInvokeFns={true}
+      />
+    );
+    $[0] = shouldReadA;
+    $[1] = a;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{ a: null, shouldReadA: true }],
+  sequentialRenders: [
+    { a: null, shouldReadA: true },
+    { a: null, shouldReadA: false },
+    { a: { b: { c: 4 } }, shouldReadA: true },
+  ],
+};
+
+```
+      
+### Eval output
+(kind: ok) [[ (exception in render) TypeError: Cannot read properties of null (reading 'b') ]]
+<div>{"fn":{"kind":"Function","result":null},"shouldInvokeFns":true}</div>
+<div>{"fn":{"kind":"Function","result":4},"shouldInvokeFns":true}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/reduce-reactive-deps/infer-function-cond-access-not-hoisted.tsx

@@ -0,0 +1,20 @@
+// @enablePropagateDepsInHIR
+
+import { Stringify } from "shared-runtime";
+
+function Foo({a, shouldReadA}) {
+  return <Stringify fn={() => {
+    if (shouldReadA) return a.b.c;
+    return null;
+  }} shouldInvokeFns={true} />
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{a: null, shouldReadA: true}],
+  sequentialRenders: [
+    {a: null, shouldReadA: true},
+    {a: null, shouldReadA: false},
+    {a: {b: {c: 4}}, shouldReadA: true},
+  ],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/reduce-reactive-deps/infer-function-uncond-access-hoisted.expect.md

@@ -0,0 +1,51 @@
+
+## Input
+
+```javascript
+// @enablePropagateDepsInHIR
+
+import { Stringify } from "shared-runtime";
+
+function useFoo(a) {
+  return <Stringify fn={() => a.b.c} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [{a: null}],
+  sequentialRenders: [{a: null}, {a: {b: {c: 4}}}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enablePropagateDepsInHIR
+
+import { Stringify } from "shared-runtime";
+
+function useFoo(a) {
+  const $ = _c(2);
+  let t0;
+  if ($[0] !== a.b.c) {
+    t0 = <Stringify fn={() => a.b.c} shouldInvokeFns={true} />;
+    $[0] = a.b.c;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [{ a: null }],
+  sequentialRenders: [{ a: null }, { a: { b: { c: 4 } } }],
+};
+
+```
+      
+### Eval output
+(kind: ok) [[ (exception in render) TypeError: Cannot read properties of undefined (reading 'c') ]]
+[[ (exception in render) TypeError: Cannot read properties of undefined (reading 'c') ]]

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/reduce-reactive-deps/infer-function-uncond-access-hoisted.tsx

@@ -0,0 +1,13 @@
+// @enablePropagateDepsInHIR
+
+import { Stringify } from "shared-runtime";
+
+function useFoo(a) {
+  return <Stringify fn={() => a.b.c} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [{a: null}],
+  sequentialRenders: [{a: null}, {a: {b: {c: 4}}}],
+};