[compiler] (needs cleanup) Hoist dependencies from functions more conservatively

Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

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

@@ -7,21 +7,27 @@ import {
   Set_intersect,
   Set_union,
   getOrInsertDefault,
+  getOrInsertWith,
 } from '../Utils/utils';
 import {
   BasicBlock,
   BlockId,
   DependencyPathEntry,
   GeneratedSource,
+  getHookKind,
   HIRFunction,
   Identifier,
   IdentifierId,
   InstructionId,
   InstructionValue,
+  LoweredFunction,
+  Place,
   PropertyLiteral,
   ReactiveScopeDependency,
   ScopeId,
+  ValidatedIdentifier,
 } from './HIR';
+import {eachInstructionOperand, eachInstructionValueOperand} from './visitors';
 
 const DEBUG_PRINT = false;
 
@@ -112,6 +118,9 @@ export function collectHoistablePropertyLoads(
     hoistableFromOptionals,
     registry,
     nestedFnImmutableContext: null,
+    fnUsages: fn.env.config.enableTreatFunctionDepsAsConditional
+      ? new Map()
+      : mapFunctionExpressionsToEscapingBlocks(fn),
   });
 }
 
@@ -127,6 +136,11 @@ type CollectHoistablePropertyLoadsContext = {
    * but are currently kept separate for readability.
    */
   nestedFnImmutableContext: ReadonlySet<IdentifierId> | null;
+  /**
+   * Mapping of functions declared within a traversal context to their
+   * (valid) usage sites, which will have hoistable property loads added
+   */
+  fnUsages: ReadonlyMap<LoweredFunction, Set<BlockId>>;
 };
 function collectHoistablePropertyLoadsImpl(
   fn: HIRFunction,
@@ -338,7 +352,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,40 +378,68 @@ function collectNonNullsInBlocks(
       ) {
         assumedNonNullObjects.add(maybeNonNull);
       }
-      if (
-        (instr.value.kind === 'FunctionExpression' ||
-          instr.value.kind === 'ObjectMethod') &&
-        !fn.env.config.enableTreatFunctionDepsAsConditional
-      ) {
+      if (instr.value.kind === 'FunctionExpression') {
+        /**
+         * What are reasonable semantics here?
+         * Risky
+         * - only treat named fns as unconditionally hoistable
+         *
+         * Conservative
+         */
         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);
+        const resultBlocks = context.fnUsages.get(innerFn);
+        if (resultBlocks != null) {
+          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),
+                ),
+              fnUsages: fn.env.config.enableTreatFunctionDepsAsConditional
+                ? new Map()
+                : mapFunctionExpressionsToEscapingBlocks(innerFn.func),
+            },
+          );
+          const innerHoistables = assertNonNull(
+            innerHoistableMap.get(innerFn.func.body.entry),
+          );
+          for (const innerBlock of resultBlocks) {
+            let innerNonNulls;
+            if (innerBlock === block.id) {
+              innerNonNulls = assumedNonNullObjects;
+            } else {
+              innerNonNulls = getOrInsertWith(nodes, innerBlock, () => ({
+                block: assertNonNull(fn.body.blocks.get(innerBlock)),
+                assumedNonNullObjects: new Set(),
+              })).assumedNonNullObjects;
+            }
+            for (const entry of innerHoistables.assumedNonNullObjects) {
+              innerNonNulls.add(entry);
+            }
+          }
         }
       }
     }
 
-    nodes.set(block.id, {
-      block,
-      assumedNonNullObjects,
-    });
+    const maybeNode = nodes.get(block.id);
+    if (maybeNode != null) {
+      // merge
+      for (const entry of assumedNonNullObjects) {
+        maybeNode.assumedNonNullObjects.add(entry);
+      }
+    } else {
+      nodes.set(block.id, {
+        block,
+        assumedNonNullObjects,
+      });
+    }
   }
   return nodes;
 }
@@ -591,3 +639,123 @@ function reduceMaybeOptionalChains(
     }
   } while (changed);
 }
+
+/**
+ *
+ * const foo = function() { ... } // this matches
+ * arr.map(function() { ... }) // this does not match
+ *
+ * What about function expressions that just escape to other functions?
+ *
+ * For both below examples, cb1 should be hoistable only to if-cond block
+ * ```js
+ * function useFoo(...) {
+ *   const cb1 = function() { ... };
+ *   const cb2 = function() { if (cond) cb1() };
+ *   return cb2;
+ * }
+ * ```
+ * ```js
+ * function useFoo(...) {
+ *   const cb1 = function() { ... };
+ *   const cb2 = function() { if (cond) return cb1; };
+ *   return cb2;
+ * }
+ * ```
+ */
+function mapFunctionExpressionsToEscapingBlocks(
+  fn: HIRFunction,
+): ReadonlyMap<LoweredFunction, Set<BlockId>> {
+  /**
+   * Step 1: gather all function expressions and known ssa'd aliases
+   */
+  const temporaries = new Map<
+    IdentifierId,
+    {fn: LoweredFunction; usage: Set<BlockId>}
+  >();
+  const validUsages = new Set<Place>();
+
+  for (const block of fn.body.blocks.values()) {
+    for (const {lvalue, value} of block.instructions) {
+      /**
+       * Only match function expressions which can have guaranteed ssa.
+       */
+      if (value.kind === 'FunctionExpression') {
+        temporaries.set(lvalue.identifier.id, {
+          fn: value.loweredFunc,
+          usage: new Set(),
+        });
+      } else if (value.kind === 'StoreLocal') {
+        const lvalue = value.lvalue.place.identifier;
+        const maybeLoweredFunc = temporaries.get(value.value.identifier.id);
+        if (
+          lvalue.name != null &&
+          lvalue.name.kind === 'named' &&
+          maybeLoweredFunc != null
+        ) {
+          temporaries.set(lvalue.id, maybeLoweredFunc);
+          validUsages.add(value.value);
+        }
+      } else if (value.kind === 'LoadLocal') {
+        const maybeLoweredFunc = temporaries.get(value.place.identifier.id);
+        if (maybeLoweredFunc != null) {
+          temporaries.set(lvalue.identifier.id, maybeLoweredFunc);
+          validUsages.add(value.place);
+        }
+      }
+    }
+  }
+  /**
+   * Step 2: Forward pass to do best-effort "escape analysis"
+   */
+  for (const block of fn.body.blocks.values()) {
+    for (const {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
+          maybeLoweredFunc.usage.add(block.id);
+        } else if (maybeHook != null) {
+          // Arguments to hooks
+          for (const arg of value.args.filter(
+            arg => arg.kind === 'Identifier',
+          ) as Array<Place>) {
+            const maybeLoweredFunc = temporaries.get(arg.identifier.id);
+            if (maybeLoweredFunc != null) {
+              maybeLoweredFunc.usage.add(block.id);
+            }
+          }
+        }
+      } else if (value.kind === 'JsxExpression') {
+        /* Match jsx attributes */
+        for (const attr of value.props) {
+          if (attr.kind === 'JsxSpreadAttribute') {
+            continue;
+          }
+          const maybeLoweredFunc = temporaries.get(attr.place.identifier.id);
+          if (maybeLoweredFunc != null) {
+            maybeLoweredFunc.usage.add(block.id);
+          }
+        }
+      }
+      if (block.terminal.kind === 'return') {
+        const maybeLoweredFunc = temporaries.get(
+          block.terminal.value.identifier.id,
+        );
+        if (maybeLoweredFunc != null) {
+          maybeLoweredFunc.usage.add(block.id);
+        }
+      }
+    }
+  }
+
+  const map = new Map<LoweredFunction, Set<BlockId>>();
+  for (const {fn, usage} of temporaries.values()) {
+    if (!map.has(fn)) {
+      map.set(fn, usage);
+    }
+  }
+  return map;
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/array-map-operand-conditionally-invoked.expect.md

@@ -0,0 +1,53 @@
+
+## Input
+
+```javascript
+function useFoo({arr}) {
+  return arr.map(e => arr[0].value + e.value);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [{arr: []}],
+  sequentialRenders: [{arr: []}, {arr: [{value: 1}, {value: 2}]}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+function useFoo(t0) {
+  const $ = _c(4);
+  const { arr } = t0;
+  let t1;
+  if ($[0] !== arr) {
+    let t2;
+    if ($[2] !== arr[0]) {
+      t2 = (e) => arr[0].value + e.value;
+      $[2] = arr[0];
+      $[3] = t2;
+    } else {
+      t2 = $[3];
+    }
+    t1 = arr.map(t2);
+    $[0] = arr;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [{ arr: [] }],
+  sequentialRenders: [{ arr: [] }, { arr: [{ value: 1 }, { value: 2 }] }],
+};
+
+```
+      
+### Eval output
+(kind: ok) []
+[2,3]

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/array-map-operand-conditionally-invoked.js

@@ -0,0 +1,9 @@
+function useFoo({arr}) {
+  return arr.map(e => arr[0].value + e.value);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [{arr: []}],
+  sequentialRenders: [{arr: []}, {arr: [{value: 1}, {value: 2}]}],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-aliased-capture-mutate.expect.md

@@ -56,9 +56,9 @@ function useFoo(t0) {
     setPropertyByKey(obj, "arr", arr);
     const obj_alias = obj;
     let t2;
-    if ($[2] !== obj_alias.arr.length) {
+    if ($[2] !== obj_alias) {
       t2 = () => obj_alias.arr.length;
-      $[2] = obj_alias.arr.length;
+      $[2] = obj_alias;
       $[3] = t2;
     } else {
       t2 = $[3];

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/preserve-memo-validation/error.maybe-invalid-useMemo-read-maybeRef.expect.md

@@ -1,32 +0,0 @@
-
-## Input
-
-```javascript
-// @validatePreserveExistingMemoizationGuarantees
-import {useMemo} from 'react';
-
-function useHook(maybeRef, shouldRead) {
-  return useMemo(() => {
-    return () => [maybeRef.current];
-  }, [shouldRead, maybeRef]);
-}
-
-```
-
-
-## Error
-
-```
-  3 |
-  4 | function useHook(maybeRef, shouldRead) {
-> 5 |   return useMemo(() => {
-    |                  ^^^^^^^
-> 6 |     return () => [maybeRef.current];
-    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-> 7 |   }, [shouldRead, maybeRef]);
-    | ^^^^ CannotPreserveMemoization: React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. The inferred dependencies did not match the manually specified dependencies, which could cause the value to change more or less frequently than expected (5:7)
-  8 | }
-  9 |
-```
-          
-      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/maybe-invalid-useMemo-read-maybeRef.expect.md

@@ -0,0 +1,40 @@
+
+## Input
+
+```javascript
+// @validatePreserveExistingMemoizationGuarantees
+import {useMemo} from 'react';
+
+function useHook(maybeRef, shouldRead) {
+  return useMemo(() => {
+    return () => [maybeRef.current];
+  }, [shouldRead, maybeRef]);
+}
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @validatePreserveExistingMemoizationGuarantees
+import { useMemo } from "react";
+
+function useHook(maybeRef, shouldRead) {
+  const $ = _c(2);
+  let t0;
+  let t1;
+  if ($[0] !== maybeRef) {
+    t1 = () => [maybeRef.current];
+    $[0] = maybeRef;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  t0 = t1;
+  return t0;
+}
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

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

@@ -41,9 +41,9 @@ function useFoo(t0) {
     local = $[1];
   }
   let t1;
-  if ($[2] !== local.b.c) {
+  if ($[2] !== local) {
     t1 = () => [() => local.b.c];
-    $[2] = local.b.c;
+    $[2] = local;
     $[3] = t1;
   } else {
     t1 = $[3];

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

@@ -34,13 +34,13 @@ function useFoo(t0) {
   const $ = _c(4);
   const { a } = t0;
   let t1;
-  if ($[0] !== a.b.c) {
+  if ($[0] !== a) {
     t1 = {
       fn() {
         return identity(a.b.c);
       },
     };
-    $[0] = a.b.c;
+    $[0] = a;
     $[1] = t1;
   } else {
     t1 = $[1];

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reactive-control-dependency-on-context-variable.expect.md

@@ -51,7 +51,7 @@ import { identity } from "shared-runtime";
 function Component(props) {
   const $ = _c(4);
   let x;
-  if ($[0] !== props.cond) {
+  if ($[0] !== props) {
     const f = () => {
       if (props.cond) {
         x = 1;
@@ -62,7 +62,7 @@ function Component(props) {
 
     const f2 = identity(f);
     f2();
-    $[0] = props.cond;
+    $[0] = props;
     $[1] = x;
   } else {
     x = $[1];

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reduce-reactive-deps/context-var-granular-dep.expect.md

@@ -82,9 +82,9 @@ function Component(t0) {
     contextVar = $[2];
   }
   let t1;
-  if ($[3] !== contextVar.val) {
+  if ($[3] !== contextVar) {
     t1 = { cb: () => contextVar.val * 4 };
-    $[3] = contextVar.val;
+    $[3] = contextVar;
     $[4] = t1;
   } else {
     t1 = $[4];

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/rename-source-variables-nested-object-method.expect.md

@@ -43,7 +43,7 @@ const t0 = "module_t0";
 const c_0 = "module_c_0";
 function useFoo(props) {
   const $0 = _c(2);
-  const c_00 = $0[0] !== props.value;
+  const c_00 = $0[0] !== props;
   let t1;
   if (c_00) {
     const a = {
@@ -61,7 +61,7 @@ function useFoo(props) {
     };
 
     t1 = a.foo().bar();
-    $0[0] = props.value;
+    $0[0] = props;
     $0[1] = t1;
   } else {
     t1 = $0[1];

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useEffect-nested-lambdas.expect.md

@@ -35,7 +35,7 @@ function Component(props) {
 import { c as _c } from "react/compiler-runtime"; // @enableTransitivelyFreezeFunctionExpressions:false
 
 function Component(props) {
-  const $ = _c(9);
+  const $ = _c(7);
   const item = useMutable(props.itemId);
   const dispatch = useDispatch();
   useFreeze(dispatch);
@@ -51,7 +51,8 @@ function Component(props) {
   }
   const exit = t0;
   let t1;
-  if ($[2] !== exit || $[3] !== item.value) {
+  let t2;
+  if ($[2] !== exit || $[3] !== item) {
     t1 = () => {
       const cleanup = GlobalEventEmitter.addListener("onInput", () => {
         if (item.value) {
@@ -60,30 +61,24 @@ function Component(props) {
       });
       return () => cleanup.remove();
     };
+    t2 = [exit, item];
     $[2] = exit;
-    $[3] = item.value;
+    $[3] = item;
     $[4] = t1;
+    $[5] = t2;
   } else {
     t1 = $[4];
-  }
-  let t2;
-  if ($[5] !== exit || $[6] !== item) {
-    t2 = [exit, item];
-    $[5] = exit;
-    $[6] = item;
-    $[7] = t2;
-  } else {
-    t2 = $[7];
+    t2 = $[5];
   }
   useEffect(t1, t2);
 
   maybeMutate(item);
   let t3;
-  if ($[8] === Symbol.for("react.memo_cache_sentinel")) {
+  if ($[6] === Symbol.for("react.memo_cache_sentinel")) {
     t3 = <div />;
-    $[8] = t3;
+    $[6] = t3;
   } else {
-    t3 = $[8];
+    t3 = $[6];
   }
   return t3;
 }