[compiler] Fix false positive for useMemo reassigning context vars (#34904)

Within a function expression local variables may use StoreContext for
local context variables, so the reassignment check here was firing too
often. We should only report an error for variables that are declared
outside the function, ie part of its `context`.

---
[//]: # (BEGIN SAPLING FOOTER)
Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with [ReviewStack](https://reviewstack.dev/facebook/react/pull/34904).
* #34903
* __->__ #34904

compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateUseMemo.ts

@@ -184,25 +184,28 @@ function validateNoContextVariableAssignment(
   fn: HIRFunction,
   errors: CompilerError,
 ): void {
+  const context = new Set(fn.context.map(place => place.identifier.id));
   for (const block of fn.body.blocks.values()) {
     for (const instr of block.instructions) {
       const value = instr.value;
       switch (value.kind) {
         case 'StoreContext': {
-          errors.pushDiagnostic(
-            CompilerDiagnostic.create({
-              category: ErrorCategory.UseMemo,
-              reason:
-                'useMemo() callbacks may not reassign variables declared outside of the callback',
-              description:
-                'useMemo() callbacks must be pure functions and cannot reassign variables defined outside of the callback function',
-              suggestions: null,
-            }).withDetails({
-              kind: 'error',
-              loc: value.lvalue.place.loc,
-              message: 'Cannot reassign variable',
-            }),
-          );
+          if (context.has(value.lvalue.place.identifier.id)) {
+            errors.pushDiagnostic(
+              CompilerDiagnostic.create({
+                category: ErrorCategory.UseMemo,
+                reason:
+                  'useMemo() callbacks may not reassign variables declared outside of the callback',
+                description:
+                  'useMemo() callbacks must be pure functions and cannot reassign variables defined outside of the callback function',
+                suggestions: null,
+              }).withDetails({
+                kind: 'error',
+                loc: value.lvalue.place.loc,
+                message: 'Cannot reassign variable',
+              }),
+            );
+          }
           break;
         }
       }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reassign-variable-in-usememo.expect.md

@@ -0,0 +1,45 @@
+
+## Input
+
+```javascript
+// @flow
+export hook useItemLanguage(items) {
+  return useMemo(() => {
+    let language: ?string = null;
+    items.forEach(item => {
+      if (item.language != null) {
+        language = item.language;
+      }
+    });
+    return language;
+  }, [items]);
+}
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+export function useItemLanguage(items) {
+  const $ = _c(2);
+  let language;
+  if ($[0] !== items) {
+    language = null;
+    items.forEach((item) => {
+      if (item.language != null) {
+        language = item.language;
+      }
+    });
+    $[0] = items;
+    $[1] = language;
+  } else {
+    language = $[1];
+  }
+  return language;
+}
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reassign-variable-in-usememo.js

@@ -0,0 +1,12 @@
+// @flow
+export hook useItemLanguage(items) {
+  return useMemo(() => {
+    let language: ?string = null;
+    items.forEach(item => {
+      if (item.language != null) {
+        language = item.language;
+      }
+    });
+    return language;
+  }, [items]);
+}