[compiler] Stop using function dependencies in propagateScopeDeps

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

@@ -231,6 +231,8 @@ const EnvironmentConfigSchema = z.object({
    */
   enableUseTypeAnnotations: z.boolean().default(false),
 
+  enableFunctionDependencyRewrite: z.boolean().default(true),
+
   /**
    * Enables inlining ReactElement object literals in place of JSX
    * An alternative to the standard JSX transform which replaces JSX with React's jsxProd() runtime

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

@@ -440,14 +440,6 @@ class Context {
 
   // Checks if identifier is a valid dependency in the current scope
   #checkValidDependency(maybeDependency: ReactiveScopeDependency): boolean {
-    // ref.current access is not a valid dep
-    if (
-      isUseRefType(maybeDependency.identifier) &&
-      maybeDependency.path.at(0)?.property === 'current'
-    ) {
-      return false;
-    }
-
     // ref value is not a valid dep
     if (isRefValueType(maybeDependency.identifier)) {
       return false;
@@ -549,6 +541,16 @@ class Context {
       });
     }
 
+    // ref.current access is not a valid dep
+    if (
+      isUseRefType(maybeDependency.identifier) &&
+      maybeDependency.path.at(0)?.property === 'current'
+    ) {
+      maybeDependency = {
+        identifier: maybeDependency.identifier,
+        path: [],
+      };
+    }
     if (this.#checkValidDependency(maybeDependency)) {
       this.#dependencies.value!.push(maybeDependency);
     }
@@ -661,35 +663,54 @@ function collectDependencies(
 
   const scopeTraversal = new ScopeBlockTraversal();
 
-  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 (!processedInstrsInOptional.has(instr)) {
-        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 (!processedInstrsInOptional.has(instr)) {
+          handleInstruction(instr, context);
+        }
       }
-    }
 
-    if (!processedInstrsInOptional.has(block.terminal)) {
-      for (const place of eachTerminalOperand(block.terminal)) {
-        context.visitOperand(place);
+      if (!processedInstrsInOptional.has(block.terminal)) {
+        for (const place of eachTerminalOperand(block.terminal)) {
+          context.visitOperand(place);
+        }
       }
     }
-  }
+  };
+
+  handleFunction(fn);
   return context.deps;
 }

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/CollectReactiveIdentifiers.ts

@@ -12,6 +12,8 @@ import {
   PrunedReactiveScopeBlock,
   ReactiveFunction,
   isPrimitiveType,
+  isUseRefType,
+  Identifier,
 } from '../HIR/HIR';
 import {ReactiveFunctionVisitor, visitReactiveFunction} from './visitors';
 
@@ -50,13 +52,21 @@ class Visitor extends ReactiveFunctionVisitor<Set<IdentifierId>> {
     this.traversePrunedScope(scopeBlock, state);
 
     for (const [id, decl] of scopeBlock.scope.declarations) {
-      if (!isPrimitiveType(decl.identifier)) {
+      if (
+        !isPrimitiveType(decl.identifier) &&
+        !isStableRefType(decl.identifier, state)
+      ) {
         state.add(id);
       }
     }
   }
 }
-
+function isStableRefType(
+  identifier: Identifier,
+  reactiveIdentifiers: Set<IdentifierId>,
+): boolean {
+  return isUseRefType(identifier) && !reactiveIdentifiers.has(identifier.id);
+}
 /*
  * Computes a set of identifiers which are reactive, using the analysis previously performed
  * in `InferReactivePlaces`.

compiler/packages/babel-plugin-react-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;
 }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/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 | }
+```
+
+