[compiler] Stop using function dependencies in propagateScopeDeps

Recursively visit inner function instructions to extract dependencies instead of using `LoweredFunction.dependencies` directly. This is currently gated by enableFunctionDependencyRewrite, which needs to be removed before we delete `LoweredFunction.dependencies` altogether (#31204). Some nice side effects: - optional-chaining deps for inner functions - full DCE and outlining for inner functions - fewer extraneous instructions
facebook · Oct 24, 2024 · c7169c0 · c7169c0
1 parent 85952e3
commit c7169c0
Show file tree

Hide file tree

Showing 15 changed files with 192 additions and 196 deletions.
diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts
@@ -230,6 +230,8 @@ const EnvironmentConfigSchema = z.object({
    */
   enableUseTypeAnnotations: z.boolean().default(false),
 
+  enableFunctionDependencyRewrite: z.boolean().default(true),
+
   /**
    * Enables inference of optional dependency chains. Without this flag
    * a property chain such as `props?.items?.foo` will infer as a dep on

diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/PropagateScopeDependenciesHIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/PropagateScopeDependenciesHIR.ts
@@ -669,35 +669,55 @@ function collectDependencies(
   ): boolean => {
     return processedInstrsInOptional.get(fn)?.has(id) ?? false;
   };
-  for (const [blockId, block] of fn.body.blocks) {
-    scopeTraversal.recordScopes(block);
-    const scopeBlockInfo = scopeTraversal.blockInfos.get(blockId);
-    if (scopeBlockInfo?.kind === 'begin') {
-      context.enterScope(scopeBlockInfo.scope);
-    } else if (scopeBlockInfo?.kind === 'end') {
-      context.exitScope(scopeBlockInfo.scope, scopeBlockInfo?.pruned);
-    }
 
-    // Record referenced optional chains in phis
-    for (const phi of block.phis) {
-      for (const operand of phi.operands) {
-        const maybeOptionalChain = temporaries.get(operand[1].identifier.id);
-        if (maybeOptionalChain) {
-          context.visitDependency(maybeOptionalChain);
+  const handleFunction = (fn: HIRFunction): void => {
+    for (const [blockId, block] of fn.body.blocks) {
+      scopeTraversal.recordScopes(block);
+      const scopeBlockInfo = scopeTraversal.blockInfos.get(blockId);
+      if (scopeBlockInfo?.kind === 'begin') {
+        context.enterScope(scopeBlockInfo.scope);
+      } else if (scopeBlockInfo?.kind === 'end') {
+        context.exitScope(scopeBlockInfo.scope, scopeBlockInfo.pruned);
+      }
+      // Record referenced optional chains in phis
+      for (const phi of block.phis) {
+        for (const operand of phi.operands) {
+          const maybeOptionalChain = temporaries.get(operand[1].identifier.id);
+          if (maybeOptionalChain) {
+            context.visitDependency(maybeOptionalChain);
+          }
         }
       }
-    }
-    for (const instr of block.instructions) {
-      if (!shouldSkipInstructionDependencies(fn, instr.id)) {
-        handleInstruction(instr, context);
+      for (const instr of block.instructions) {
+        if (
+          fn.env.config.enableFunctionDependencyRewrite &&
+          (instr.value.kind === 'FunctionExpression' ||
+            instr.value.kind === 'ObjectMethod')
+        ) {
+          context.declare(instr.lvalue.identifier, {
+            id: instr.id,
+            scope: context.currentScope,
+          });
+          /**
+           * Recursively visit the inner function to extract dependencies there
+           */
+          const wasInInnerFn = context.inInnerFn;
+          context.inInnerFn = true;
+          handleFunction(instr.value.loweredFunc.func);
+          context.inInnerFn = wasInInnerFn;
+        } else if (!shouldSkipInstructionDependencies(fn, instr.id)) {
+          handleInstruction(instr, context);
+        }
       }
-    }
 
-    if (!shouldSkipInstructionDependencies(fn, block.terminal.id)) {
-      for (const place of eachTerminalOperand(block.terminal)) {
-        context.visitOperand(place);
+      if (!shouldSkipInstructionDependencies(fn, block.terminal.id)) {
+        for (const place of eachTerminalOperand(block.terminal)) {
+          context.visitOperand(place);
+        }
       }
     }
-  }
+  };
+
+  handleFunction(fn);
   return context.deps;
 }
diff --git a/...eact-compiler/src/__tests__/fixtures/compiler/capturing-func-mutate-2.expect.md b/...eact-compiler/src/__tests__/fixtures/compiler/capturing-func-mutate-2.expect.md
@@ -26,29 +26,20 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { c as _c } from "react/compiler-runtime";
 function component(a, b) {
-  const $ = _c(5);
-  let t0;
-  if ($[0] !== b) {
-    t0 = { b };
-    $[0] = b;
-    $[1] = t0;
-  } else {
-    t0 = $[1];
-  }
-  const y = t0;
+  const $ = _c(2);
+  const y = { b };
   let z;
-  if ($[2] !== a || $[3] !== y) {
+  if ($[0] !== a) {
     z = { a };
     const x = function () {
       z.a = 2;
     };
 
     x();
-    $[2] = a;
-    $[3] = y;
-    $[4] = z;
+    $[0] = a;
+    $[1] = z;
   } else {
-    z = $[4];
+    z = $[1];
   }
   return z;
 }

diff --git a/...iler/src/__tests__/fixtures/compiler/jsx-outlining-child-stored-in-id.expect.md b/...iler/src/__tests__/fixtures/compiler/jsx-outlining-child-stored-in-id.expect.md
@@ -53,7 +53,7 @@ function Component(arr) {
   const $ = _c(3);
   const x = useX();
   let t0;
-  if ($[0] !== arr || $[1] !== x) {
+  if ($[0] !== x || $[1] !== arr) {
     t0 = arr.map((i) => {
       arr.map((i_0, id) => {
         const T0 = _temp;
@@ -63,8 +63,8 @@ function Component(arr) {
         return jsx;
       });
     });
-    $[0] = arr;
-    $[1] = x;
+    $[0] = x;
+    $[1] = arr;
     $[2] = t0;
   } else {
     t0 = $[2];

diff --git a/...alidation/error.todo-useCallback-captures-reassigned-context-property.expect.md b/...alidation/error.todo-useCallback-captures-reassigned-context-property.expect.md
@@ -0,0 +1,53 @@
+
+## Input
+
+```javascript
+// @validatePreserveExistingMemoizationGuarantees
+import {useCallback} from 'react';
+import {Stringify} from 'shared-runtime';
+
+/**
+ * TODO: we're currently bailing out because `contextVar` is a context variable
+ * and not recorded into the PropagateScopeDeps LoadLocal / PropertyLoad
+ * sidemap. Previously, we were able to avoid this as `BuildHIR` hoisted
+ * `LoadContext` and `PropertyLoad` instructions into the outer function, which
+ * we took as eligible dependencies.
+ *
+ * One solution is to simply record `LoadContext` identifiers into the
+ * temporaries sidemap when the instruction occurs *after* the context
+ * variable's mutable range.
+ */
+function Foo(props) {
+  let contextVar;
+  if (props.cond) {
+    contextVar = {val: 2};
+  } else {
+    contextVar = {};
+  }
+
+  const cb = useCallback(() => [contextVar.val], [contextVar.val]);
+
+  return <Stringify cb={cb} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{cond: true}],
+};
+
+```
+
+
+## Error
+
+```
+  22 |   }
+  23 |
+> 24 |   const cb = useCallback(() => [contextVar.val], [contextVar.val]);
+     |                          ^^^^^^^^^^^^^^^^^^^^^^ 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 (24:24)
+  25 |
+  26 |   return <Stringify cb={cb} shouldInvokeFns={true} />;
+  27 | }
+```
+          
+      
diff --git a/.../preserve-memo-validation/error.todo-useCallback-captures-reassigned-context-property.tsx b/.../preserve-memo-validation/error.todo-useCallback-captures-reassigned-context-property.tsx
@@ -0,0 +1,32 @@
+// @validatePreserveExistingMemoizationGuarantees
+import {useCallback} from 'react';
+import {Stringify} from 'shared-runtime';
+
+/**
+ * TODO: we're currently bailing out because `contextVar` is a context variable
+ * and not recorded into the PropagateScopeDeps LoadLocal / PropertyLoad
+ * sidemap. Previously, we were able to avoid this as `BuildHIR` hoisted
+ * `LoadContext` and `PropertyLoad` instructions into the outer function, which
+ * we took as eligible dependencies.
+ *
+ * One solution is to simply record `LoadContext` identifiers into the
+ * temporaries sidemap when the instruction occurs *after* the context
+ * variable's mutable range.
+ */
+function Foo(props) {
+  let contextVar;
+  if (props.cond) {
+    contextVar = {val: 2};
+  } else {
+    contextVar = {};
+  }
+
+  const cb = useCallback(() => [contextVar.val], [contextVar.val]);
+
+  return <Stringify cb={cb} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{cond: true}],
+};
diff --git a/...o-validation/error.useCallback-infer-less-specific-conditional-access.expect.md b/...o-validation/error.useCallback-infer-less-specific-conditional-access.expect.md
@@ -44,8 +44,6 @@ function Component({propA, propB}) {
      | ^^^^^^^^^^^^^^^^^
 > 14 |   }, [propA?.a, propB.x.y]);
      | ^^^^ 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 (6:14)
-
-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 (6:14)
   15 | }
   16 |
 ```

diff --git a/...erve-memo-validation/useCallback-captures-reassigned-context-property.expect.md b/...erve-memo-validation/useCallback-captures-reassigned-context-property.expect.md
diff --git a/...es/compiler/preserve-memo-validation/useCallback-captures-reassigned-context-property.tsx b/...es/compiler/preserve-memo-validation/useCallback-captures-reassigned-context-property.tsx
diff --git a/...iler/preserve-memo-validation/useCallback-captures-reassigned-context.expect.md b/...iler/preserve-memo-validation/useCallback-captures-reassigned-context.expect.md
@@ -45,18 +45,16 @@ function Foo(props) {
   } else {
     x = $[1];
   }
-
-  const t0 = x;
-  let t1;
-  if ($[2] !== t0) {
-    t1 = () => [x];
-    $[2] = t0;
-    $[3] = t1;
+  let t0;
+  if ($[2] !== x) {
+    t0 = () => [x];
+    $[2] = x;
+    $[3] = t0;
   } else {
-    t1 = $[3];
+    t0 = $[3];
   }
   x;
-  const cb = t1;
+  const cb = t0;
   return cb;
 }