[compiler][be] Clean up nested function context in DCE (#31202)

Now that we rely on function context exclusively, let's clean up
`HIRFunction.context` after DCE. This PR is in preparation of #31204,
which would otherwise have unnecessary declarations (of context values
that become entirely DCE'd)

'
---
[//]: # (BEGIN SAPLING FOOTER)
Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with [ReviewStack](https://reviewstack.dev/facebook/react/pull/31202).
* __->__ #31202
* #31203
* #31201
* #31200
* #31521

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/__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 });
+}
 
 ```
       

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/function-expr-directive.expect.md

@@ -34,7 +34,7 @@ function Component() {
     t0 = function update() {
       "worklet";
 
-      setCount((count_0) => count_0 + 1);
+      setCount(_temp);
     };
     $[0] = t0;
   } else {
@@ -51,6 +51,9 @@ function Component() {
   }
   return t1;
 }
+function _temp(count_0) {
+  return count_0 + 1;
+}
 
 export const FIXTURE_ENTRYPOINT = {
   fn: Component,

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/merge-scopes-callback.expect.md

@@ -32,7 +32,7 @@ function Component() {
   let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
     t0 = () => {
-      setState((s) => s + 1);
+      setState(_temp);
     };
     $[0] = t0;
   } else {
@@ -61,6 +61,9 @@ function Component() {
   }
   return t2;
 }
+function _temp(s) {
+  return s + 1;
+}
 
 ```
       

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-no-declarations-in-reactive-scope-with-early-return.expect.md

@@ -39,7 +39,7 @@ function Component() {
 ```javascript
 import { c as _c } from "react/compiler-runtime"; // @enableAssumeHooksFollowRulesOfReact @enableTransitivelyFreezeFunctionExpressions
 function Component() {
-  const $ = _c(8);
+  const $ = _c(7);
   const items = useItems();
   let t0;
   let t1;
@@ -47,35 +47,25 @@ function Component() {
   if ($[0] !== items) {
     t2 = Symbol.for("react.early_return_sentinel");
     bb0: {
-      let t3;
-      if ($[4] === Symbol.for("react.memo_cache_sentinel")) {
-        t3 = (t4) => {
-          const [item] = t4;
-          return item.name != null;
-        };
-        $[4] = t3;
-      } else {
-        t3 = $[4];
-      }
-      t0 = items.filter(t3);
+      t0 = items.filter(_temp);
       const filteredItems = t0;
       if (filteredItems.length === 0) {
-        let t4;
-        if ($[5] === Symbol.for("react.memo_cache_sentinel")) {
-          t4 = (
+        let t3;
+        if ($[4] === Symbol.for("react.memo_cache_sentinel")) {
+          t3 = (
             <div>
               <span />
             </div>
           );
-          $[5] = t4;
+          $[4] = t3;
         } else {
-          t4 = $[5];
+          t3 = $[4];
         }
-        t2 = t4;
+        t2 = t3;
         break bb0;
       }
 
-      t1 = filteredItems.map(_temp);
+      t1 = filteredItems.map(_temp2);
     }
     $[0] = items;
     $[1] = t1;
@@ -90,19 +80,23 @@ function Component() {
     return t2;
   }
   let t3;
-  if ($[6] !== t1) {
+  if ($[5] !== t1) {
     t3 = <>{t1}</>;
-    $[6] = t1;
-    $[7] = t3;
+    $[5] = t1;
+    $[6] = t3;
   } else {
-    t3 = $[7];
+    t3 = $[6];
   }
   return t3;
 }
-function _temp(t0) {
+function _temp2(t0) {
   const [item_0] = t0;
   return <Stringify item={item_0} />;
 }
+function _temp(t0) {
+  const [item] = t0;
+  return item.name != null;
+}
 
 ```
       

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/resolve-react-hooks-based-on-import-name.expect.md

@@ -38,7 +38,7 @@ function Component() {
   let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
     t0 = () => {
-      setState((s) => s + 1);
+      setState(_temp);
     };
     $[0] = t0;
   } else {
@@ -67,6 +67,9 @@ function Component() {
   }
   return t2;
 }
+function _temp(s) {
+  return s + 1;
+}
 
 export const FIXTURE_ENTRYPOINT = {
   fn: Component,