From b6b33bfb92c095160df7370fb488acb89c55b5ca Mon Sep 17 00:00:00 2001
From: mofeiZ <34200447+mofeiZ@users.noreply.github.com>
Date: Wed, 22 Jan 2025 14:34:08 -0500
Subject: [PATCH] [compiler][ez] rewrite invariant in InferReferenceEffects
 (#32093)

Small patch to pass aliased context values into
`Object|ArrayExpression`s
---
[//]: # (BEGIN SAPLING FOOTER)
Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with [ReviewStack](https://reviewstack.dev/facebook/react/pull/32093).
* #32099
* #32104
* #32098
* #32097
* #32096
* #32095
* #32094
* __->__ #32093
---
 .../src/Inference/InferFunctionEffects.ts     |  7 +-
 .../src/Inference/InferReferenceEffects.ts    | 75 +++++++++++--------
 2 files changed, 49 insertions(+), 33 deletions(-)

diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferFunctionEffects.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferFunctionEffects.ts
index 0ae54839b6fa3..fe209924e4e08 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferFunctionEffects.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferFunctionEffects.ts
@@ -41,11 +41,16 @@ function inferOperandEffect(state: State, place: Place): null | FunctionEffect {
       if (isRefOrRefValue(place.identifier)) {
         break;
       } else if (value.kind === ValueKind.Context) {
+        CompilerError.invariant(value.context.size > 0, {
+          reason:
+            "[InferFunctionEffects] Expected Context-kind value's capture list to be non-empty.",
+          loc: place.loc,
+        });
         return {
           kind: 'ContextMutation',
           loc: place.loc,
           effect: place.effect,
-          places: value.context.size === 0 ? new Set([place]) : value.context,
+          places: value.context,
         };
       } else if (
         value.kind !== ValueKind.Mutable &&
diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts
index 8cf30a9666e25..70395e58d930e 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts
@@ -857,17 +857,19 @@ function inferBlock(
         break;
       }
       case 'ArrayExpression': {
-        const valueKind: AbstractValue = hasContextRefOperand(state, instrValue)
-          ? {
-              kind: ValueKind.Context,
-              reason: new Set([ValueReason.Other]),
-              context: new Set(),
-            }
-          : {
-              kind: ValueKind.Mutable,
-              reason: new Set([ValueReason.Other]),
-              context: new Set(),
-            };
+        const contextRefOperands = getContextRefOperand(state, instrValue);
+        const valueKind: AbstractValue =
+          contextRefOperands.length > 0
+            ? {
+                kind: ValueKind.Context,
+                reason: new Set([ValueReason.Other]),
+                context: new Set(contextRefOperands),
+              }
+            : {
+                kind: ValueKind.Mutable,
+                reason: new Set([ValueReason.Other]),
+                context: new Set(),
+              };
         continuation = {
           kind: 'initialize',
           valueKind,
@@ -918,17 +920,19 @@ function inferBlock(
         break;
       }
       case 'ObjectExpression': {
-        const valueKind: AbstractValue = hasContextRefOperand(state, instrValue)
-          ? {
-              kind: ValueKind.Context,
-              reason: new Set([ValueReason.Other]),
-              context: new Set(),
-            }
-          : {
-              kind: ValueKind.Mutable,
-              reason: new Set([ValueReason.Other]),
-              context: new Set(),
-            };
+        const contextRefOperands = getContextRefOperand(state, instrValue);
+        const valueKind: AbstractValue =
+          contextRefOperands.length > 0
+            ? {
+                kind: ValueKind.Context,
+                reason: new Set([ValueReason.Other]),
+                context: new Set(contextRefOperands),
+              }
+            : {
+                kind: ValueKind.Mutable,
+                reason: new Set([ValueReason.Other]),
+                context: new Set(),
+              };
 
         for (const property of instrValue.properties) {
           switch (property.kind) {
@@ -1593,15 +1597,21 @@ function inferBlock(
       }
       case 'LoadLocal': {
         const lvalue = instr.lvalue;
-        const effect =
-          state.isDefined(lvalue) &&
-          state.kind(lvalue).kind === ValueKind.Context
-            ? Effect.ConditionallyMutate
-            : Effect.Capture;
+        CompilerError.invariant(
+          !(
+            state.isDefined(lvalue) &&
+            state.kind(lvalue).kind === ValueKind.Context
+          ),
+          {
+            reason:
+              '[InferReferenceEffects] Unexpected LoadLocal with context kind',
+            loc: lvalue.loc,
+          },
+        );
         state.referenceAndRecordEffects(
           freezeActions,
           instrValue.place,
-          effect,
+          Effect.Capture,
           ValueReason.Other,
         );
         lvalue.effect = Effect.ConditionallyMutate;
@@ -1932,19 +1942,20 @@ function inferBlock(
   );
 }
 
-function hasContextRefOperand(
+function getContextRefOperand(
   state: InferenceState,
   instrValue: InstructionValue,
-): boolean {
+): Array<Place> {
+  const result = [];
   for (const place of eachInstructionValueOperand(instrValue)) {
     if (
       state.isDefined(place) &&
       state.kind(place).kind === ValueKind.Context
     ) {
-      return true;
+      result.push(place);
     }
   }
-  return false;
+  return result;
 }
 
 export function getFunctionCallSignature(