[compiler][be] Clean up nested function context in DCE

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

@@ -3186,7 +3186,13 @@ function lowerJsxMemberExpression(
       loc: object.node.loc ?? null,
       suggestions: null,
     });
-    objectPlace = lowerIdentifier(builder, object);
+
+    const kind = getLoadKind(builder, object);
+    objectPlace = lowerValueToTemporary(builder, {
+      kind: kind,
+      place: lowerIdentifier(builder, object),
+      loc: exprPath.node.loc ?? GeneratedSource,
+    });
   }
   const property = exprPath.get('property').node.name;
   return lowerValueToTemporary(builder, {

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/Optimization/DeadCodeElimination.ts

@@ -58,6 +58,14 @@ export function deadCodeElimination(fn: HIRFunction): void {
       }
     }
   }
+
+  /**
+   * Constant propagation and DCE may have deleted or rewritten instructions
+   * that reference context variables.
+   */
+  retainWhere(fn.context, contextVar =>
+    state.isIdOrNameUsed(contextVar.identifier),
+  );
 }
 
 class State {

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/arrow-expr-directive.expect.md

@@ -28,7 +28,7 @@ function Component() {
     t0 = () => {
       "worklet";
 
-      setCount((count_0) => count_0 + 1);
+      setCount(_temp);
     };
     $[0] = t0;
   } else {
@@ -45,6 +45,9 @@ function Component() {
   }
   return t1;
 }
+function _temp(count_0) {
+  return count_0 + 1;
+}
 
 ```
       

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

@@ -55,11 +55,7 @@ function getNativeLogFunction(level) {
       if (arguments.length === 1 && typeof arguments[0] === "string") {
         str = arguments[0];
       } else {
-        str = Array.prototype.map
-          .call(arguments, function (arg) {
-            return inspect(arg, { depth: 10 });
-          })
-          .join(", ");
+        str = Array.prototype.map.call(arguments, _temp).join(", ");
       }
 
       const firstArg = arguments[0];
@@ -92,6 +88,9 @@ function getNativeLogFunction(level) {
   }
   return t0;
 }
+function _temp(arg) {
+  return inspect(arg, { depth: 10 });
+}
 
 ```