[compiler] Hoist dependencies from functions more conservatively

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

@@ -13,11 +13,13 @@ import {
   BlockId,
   DependencyPathEntry,
   GeneratedSource,
+  getHookKind,
   HIRFunction,
   Identifier,
   IdentifierId,
   InstructionId,
   InstructionValue,
+  LoweredFunction,
   PropertyLiteral,
   ReactiveScopeDependency,
   ScopeId,
@@ -112,6 +114,9 @@ export function collectHoistablePropertyLoads(
     hoistableFromOptionals,
     registry,
     nestedFnImmutableContext: null,
+    assumedInvokedFns: fn.env.config.enableTreatFunctionDepsAsConditional
+      ? new Set()
+      : getAssumedInvokedFunctions(fn),
   });
 }
 
@@ -127,6 +132,11 @@ type CollectHoistablePropertyLoadsContext = {
    * but are currently kept separate for readability.
    */
   nestedFnImmutableContext: ReadonlySet<IdentifierId> | null;
+  /**
+   * Functions which are assumed to be eventually called (as opposed to ones which might
+   * not be called, e.g. the 0th argument of Array.map)
+   */
+  assumedInvokedFns: ReadonlySet<LoweredFunction>;
 };
 function collectHoistablePropertyLoadsImpl(
   fn: HIRFunction,
@@ -338,7 +348,13 @@ function collectNonNullsInBlocks(
       context.registry.getOrCreateIdentifier(identifier),
     );
   }
-  const nodes = new Map<BlockId, BlockInfo>();
+  const nodes = new Map<
+    BlockId,
+    {
+      block: BasicBlock;
+      assumedNonNullObjects: Set<PropertyPathNode>;
+    }
+  >();
   for (const [_, block] of fn.body.blocks) {
     const assumedNonNullObjects = new Set<PropertyPathNode>(
       knownNonNullIdentifiers,
@@ -358,32 +374,30 @@ function collectNonNullsInBlocks(
       ) {
         assumedNonNullObjects.add(maybeNonNull);
       }
-      if (
-        (instr.value.kind === 'FunctionExpression' ||
-          instr.value.kind === 'ObjectMethod') &&
-        !fn.env.config.enableTreatFunctionDepsAsConditional
-      ) {
+      if (instr.value.kind === 'FunctionExpression') {
         const innerFn = instr.value.loweredFunc;
-        const innerHoistableMap = collectHoistablePropertyLoadsImpl(
-          innerFn.func,
-          {
-            ...context,
-            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),
-        );
-        for (const entry of innerHoistables.assumedNonNullObjects) {
-          assumedNonNullObjects.add(entry);
+        if (context.assumedInvokedFns.has(innerFn)) {
+          const innerHoistableMap = collectHoistablePropertyLoadsImpl(
+            innerFn.func,
+            {
+              ...context,
+              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),
+          );
+          for (const entry of innerHoistables.assumedNonNullObjects) {
+            assumedNonNullObjects.add(entry);
+          }
         }
       }
     }
@@ -591,3 +605,130 @@ function reduceMaybeOptionalChains(
     }
   } while (changed);
 }
+
+function getAssumedInvokedFunctions(
+  fn: HIRFunction,
+  temporaries: Map<
+    IdentifierId,
+    {fn: LoweredFunction; mayInvoke: Set<LoweredFunction>}
+  > = new Map(),
+): ReadonlySet<LoweredFunction> {
+  const hoistableFunctions = new Set<LoweredFunction>();
+  /**
+   * Step 1: Conservatively collect identifier to function expression mappings
+   */
+  for (const block of fn.body.blocks.values()) {
+    for (const {lvalue, value} of block.instructions) {
+      /**
+       * Conservatively only match function expressions which can have guaranteed ssa.
+       * ObjectMethods and ObjectProperties do not.
+       */
+      if (value.kind === 'FunctionExpression') {
+        temporaries.set(lvalue.identifier.id, {
+          fn: value.loweredFunc,
+          mayInvoke: new Set(),
+        });
+      } else if (value.kind === 'StoreLocal') {
+        const lvalue = value.lvalue.place.identifier;
+        const maybeLoweredFunc = temporaries.get(value.value.identifier.id);
+        if (maybeLoweredFunc != null) {
+          temporaries.set(lvalue.id, maybeLoweredFunc);
+        }
+      } else if (value.kind === 'LoadLocal') {
+        const maybeLoweredFunc = temporaries.get(value.place.identifier.id);
+        if (maybeLoweredFunc != null) {
+          temporaries.set(lvalue.identifier.id, maybeLoweredFunc);
+        }
+      }
+    }
+  }
+  /**
+   * Step 2: Forward pass to do analysis of assumed function calls. Note that
+   * this is conservative and does not count indirect references through
+   * containers (e.g. `return {cb: () => {...}})`).
+   */
+  for (const block of fn.body.blocks.values()) {
+    for (const {lvalue, value} of block.instructions) {
+      if (value.kind === 'CallExpression') {
+        const callee = value.callee;
+        const maybeHook = getHookKind(fn.env, callee.identifier);
+        const maybeLoweredFunc = temporaries.get(callee.identifier.id);
+        if (maybeLoweredFunc != null) {
+          // Direct calls
+          hoistableFunctions.add(maybeLoweredFunc.fn);
+        } else if (maybeHook != null) {
+          /**
+           * Assume arguments to all hooks are safe to invoke
+           */
+          for (const arg of value.args) {
+            if (arg.kind === 'Identifier') {
+              const maybeLoweredFunc = temporaries.get(arg.identifier.id);
+              if (maybeLoweredFunc != null) {
+                hoistableFunctions.add(maybeLoweredFunc.fn);
+              }
+            }
+          }
+        }
+      } else if (value.kind === 'JsxExpression') {
+        /**
+         * Assume JSX attributes and children are safe to invoke
+         */
+        for (const attr of value.props) {
+          if (attr.kind === 'JsxSpreadAttribute') {
+            continue;
+          }
+          const maybeLoweredFunc = temporaries.get(attr.place.identifier.id);
+          if (maybeLoweredFunc != null) {
+            hoistableFunctions.add(maybeLoweredFunc.fn);
+          }
+        }
+        for (const child of value.children ?? []) {
+          const maybeLoweredFunc = temporaries.get(child.identifier.id);
+          if (maybeLoweredFunc != null) {
+            hoistableFunctions.add(maybeLoweredFunc.fn);
+          }
+        }
+      } else if (value.kind === 'FunctionExpression') {
+        /**
+         * Recursively traverse into other function expressions which may invoke
+         * or pass already declared functions to react (e.g. as JSXAttributes).
+         *
+         * If lambda A calls lambda B, we assume lambda B is safe to invoke if
+         * lambda A is -- even if lambda B is conditionally called. (see
+         * `conditional-call-chain` fixture for example).
+         */
+        const loweredFunc = value.loweredFunc.func;
+        const lambdasCalled = getAssumedInvokedFunctions(
+          loweredFunc,
+          temporaries,
+        );
+        const maybeLoweredFunc = temporaries.get(lvalue.identifier.id);
+        if (maybeLoweredFunc != null) {
+          for (const called of lambdasCalled) {
+            maybeLoweredFunc.mayInvoke.add(called);
+          }
+        }
+      }
+    }
+    if (block.terminal.kind === 'return') {
+      /**
+       * Assume directly returned functions are safe to call
+       */
+      const maybeLoweredFunc = temporaries.get(
+        block.terminal.value.identifier.id,
+      );
+      if (maybeLoweredFunc != null) {
+        hoistableFunctions.add(maybeLoweredFunc.fn);
+      }
+    }
+  }
+
+  for (const [_, {fn, mayInvoke}] of temporaries) {
+    if (hoistableFunctions.has(fn)) {
+      for (const called of mayInvoke) {
+        hoistableFunctions.add(called);
+      }
+    }
+  }
+  return hoistableFunctions;
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/function-expression-prototype-call.expect.md

@@ -23,11 +23,11 @@ import { c as _c } from "react/compiler-runtime";
 function Component(props) {
   const $ = _c(4);
   let t0;
-  if ($[0] !== props.name) {
+  if ($[0] !== props) {
     t0 = function () {
       return <div>{props.name}</div>;
     };
-    $[0] = props.name;
+    $[0] = props;
     $[1] = t0;
   } else {
     t0 = $[1];

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inner-function/nullable-objects/array-map-named-callback-cross-context.expect.md

@@ -0,0 +1,133 @@
+
+## Input
+
+```javascript
+import {Stringify} from 'shared-runtime';
+
+/**
+ * Forked from array-map-simple.js
+ *
+ * Named lambdas (e.g. cb1) may be defined in the top scope of a function and
+ * used in a different lambda (getArrMap1).
+ *
+ * Here, we should try to determine if cb1 is actually called. In this case:
+ * - getArrMap1 is assumed to be called as it's passed to JSX
+ * - cb1 is not assumed to be called since it's only used as a call operand
+ */
+function useFoo({arr1, arr2}) {
+  const cb1 = e => arr1[0].value + e.value;
+  const getArrMap1 = () => arr1.map(cb1);
+  const cb2 = e => arr2[0].value + e.value;
+  const getArrMap2 = () => arr1.map(cb2);
+  return (
+    <Stringify
+      getArrMap1={getArrMap1}
+      getArrMap2={getArrMap2}
+      shouldInvokeFns={true}
+    />
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [{arr1: [], arr2: []}],
+  sequentialRenders: [
+    {arr1: [], arr2: []},
+    {arr1: [], arr2: null},
+    {arr1: [{value: 1}, {value: 2}], arr2: [{value: -1}]},
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { Stringify } from "shared-runtime";
+
+/**
+ * Forked from array-map-simple.js
+ *
+ * Named lambdas (e.g. cb1) may be defined in the top scope of a function and
+ * used in a different lambda (getArrMap1).
+ *
+ * Here, we should try to determine if cb1 is actually called. In this case:
+ * - getArrMap1 is assumed to be called as it's passed to JSX
+ * - cb1 is not assumed to be called since it's only used as a call operand
+ */
+function useFoo(t0) {
+  const $ = _c(13);
+  const { arr1, arr2 } = t0;
+  let t1;
+  if ($[0] !== arr1[0]) {
+    t1 = (e) => arr1[0].value + e.value;
+    $[0] = arr1[0];
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  const cb1 = t1;
+  let t2;
+  if ($[2] !== arr1 || $[3] !== cb1) {
+    t2 = () => arr1.map(cb1);
+    $[2] = arr1;
+    $[3] = cb1;
+    $[4] = t2;
+  } else {
+    t2 = $[4];
+  }
+  const getArrMap1 = t2;
+  let t3;
+  if ($[5] !== arr2) {
+    t3 = (e_0) => arr2[0].value + e_0.value;
+    $[5] = arr2;
+    $[6] = t3;
+  } else {
+    t3 = $[6];
+  }
+  const cb2 = t3;
+  let t4;
+  if ($[7] !== arr1 || $[8] !== cb2) {
+    t4 = () => arr1.map(cb2);
+    $[7] = arr1;
+    $[8] = cb2;
+    $[9] = t4;
+  } else {
+    t4 = $[9];
+  }
+  const getArrMap2 = t4;
+  let t5;
+  if ($[10] !== getArrMap1 || $[11] !== getArrMap2) {
+    t5 = (
+      <Stringify
+        getArrMap1={getArrMap1}
+        getArrMap2={getArrMap2}
+        shouldInvokeFns={true}
+      />
+    );
+    $[10] = getArrMap1;
+    $[11] = getArrMap2;
+    $[12] = t5;
+  } else {
+    t5 = $[12];
+  }
+  return t5;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [{ arr1: [], arr2: [] }],
+  sequentialRenders: [
+    { arr1: [], arr2: [] },
+    { arr1: [], arr2: null },
+    { arr1: [{ value: 1 }, { value: 2 }], arr2: [{ value: -1 }] },
+  ],
+};
+
+```
+
+### Eval output
+(kind: ok) <div>{"getArrMap1":{"kind":"Function","result":[]},"getArrMap2":{"kind":"Function","result":[]},"shouldInvokeFns":true}</div>
+<div>{"getArrMap1":{"kind":"Function","result":[]},"getArrMap2":{"kind":"Function","result":[]},"shouldInvokeFns":true}</div>
+<div>{"getArrMap1":{"kind":"Function","result":[2,3]},"getArrMap2":{"kind":"Function","result":[0,1]},"shouldInvokeFns":true}</div>