Update base for Update on "[compiler] Separate InferFunctionEffects pass from InferReferenceEffects"

Test Plan:
This diff finally separates InferFunctionEffects into its own separate pass from InferReferenceEffects. It relies on the abstractValues populated in IRE as well as the alias sets that IRE returns.

The meat of the InferFunctionEffects algorithm is still the same, but rather than querying the "live" InferenceState from IRE for abstract values and alias information, we query the values defined on places and the alias set returned by IRE.

One extra bit of work that we now need to perform is creating a map from IdentifierIds to AbstractValues, which we do by traversing the HIRFunction and examining all places. We also need to track the computed effects of nested functions, and that's where the disjoint set of aliases comes in -- when performing the main algorithm, we might see a set of instructions like

```
$0 = Function (effect=ContextMutation) { ... }
$1 = LoadLocal $0
$2 = Call $1 ()
```

When examining the `Call` instruction, we need to know that $1 has the function effect [ContextMutation]. Since $0 and $1 were inferred to be aliased by IRE, we don't need to do any other propagation from $0 to $1 if we track the nested effects based on the root of the alias set { $0, $1 } rather than the specific identifier $1 present in the Call instruction.

[ghstack-poisoned]

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

@@ -82,6 +82,7 @@ export function lower(
       kind: 'Identifier',
       identifier: builder.resolveBinding(ref),
       effect: Effect.Unknown,
+      abstractValue: null,
       reactive: false,
       loc: ref.loc ?? GeneratedSource,
     });
@@ -113,6 +114,7 @@ export function lower(
         kind: 'Identifier',
         identifier: binding.identifier,
         effect: Effect.Unknown,
+        abstractValue: null,
         reactive: false,
         loc: param.node.loc ?? GeneratedSource,
       };
@@ -126,6 +128,7 @@ export function lower(
         kind: 'Identifier',
         identifier: builder.makeTemporary(param.node.loc ?? GeneratedSource),
         effect: Effect.Unknown,
+        abstractValue: null,
         reactive: false,
         loc: param.node.loc ?? GeneratedSource,
       };
@@ -144,6 +147,7 @@ export function lower(
         kind: 'Identifier',
         identifier: builder.makeTemporary(param.node.loc ?? GeneratedSource),
         effect: Effect.Unknown,
+        abstractValue: null,
         reactive: false,
         loc: param.node.loc ?? GeneratedSource,
       };
@@ -460,6 +464,7 @@ function lowerStatement(
           });
           const place: Place = {
             effect: Effect.Unknown,
+            abstractValue: null,
             identifier: identifier.identifier,
             kind: 'Identifier',
             reactive: false,
@@ -853,6 +858,7 @@ function lowerStatement(
           } else {
             const place: Place = {
               effect: Effect.Unknown,
+              abstractValue: null,
               identifier: binding.identifier,
               kind: 'Identifier',
               reactive: false,
@@ -1264,6 +1270,7 @@ function lowerStatement(
             handlerBindingPath.node.loc ?? GeneratedSource,
           ),
           effect: Effect.Unknown,
+          abstractValue: null,
           reactive: false,
           loc: handlerBindingPath.node.loc ?? GeneratedSource,
         };
@@ -3428,6 +3435,7 @@ function lowerIdentifier(
         kind: 'Identifier',
         identifier: binding.identifier,
         effect: Effect.Unknown,
+        abstractValue: null,
         reactive: false,
         loc: exprLoc,
       };
@@ -3449,6 +3457,7 @@ function buildTemporaryPlace(builder: HIRBuilder, loc: SourceLocation): Place {
     kind: 'Identifier',
     identifier: builder.makeTemporary(loc),
     effect: Effect.Unknown,
+    abstractValue: null,
     reactive: false,
     loc,
   };
@@ -3511,6 +3520,7 @@ function lowerIdentifierForAssignment(
     kind: 'Identifier',
     identifier: binding.identifier,
     effect: Effect.Unknown,
+    abstractValue: null,
     reactive: false,
     loc,
   };

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

@@ -761,6 +761,7 @@ function _staticInvariantInstructionValueHasLocation(
 export type Phi = {
   kind: 'Phi';
   id: Identifier;
+  abstractValue: AbstractValue | null;
   operands: Map<BlockId, Identifier>;
 };
 
@@ -1110,6 +1111,7 @@ export type Place = {
   kind: 'Identifier';
   identifier: Identifier;
   effect: Effect;
+  abstractValue: AbstractValue | null;
   reactive: boolean;
   loc: SourceLocation;
 };

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

@@ -895,6 +895,7 @@ export function createTemporaryPlace(
     kind: 'Identifier',
     identifier: makeTemporaryIdentifier(env.nextIdentifierId, loc),
     reactive: false,
+    abstractValue: null,
     effect: Effect.Unknown,
     loc: GeneratedSource,
   };

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

@@ -86,6 +86,7 @@ export function mergeConsecutiveBlocks(fn: HIRFunction): void {
           kind: 'Identifier',
           identifier: phi.id,
           effect: Effect.ConditionallyMutate,
+          abstractValue: null,
           reactive: false,
           loc: GeneratedSource,
         },
@@ -95,6 +96,7 @@ export function mergeConsecutiveBlocks(fn: HIRFunction): void {
             kind: 'Identifier',
             identifier: operand,
             effect: Effect.Read,
+            abstractValue: null,
             reactive: false,
             loc: GeneratedSource,
           },

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

@@ -833,6 +833,8 @@ export function printPattern(pattern: Pattern | Place | SpreadPattern): string {
 
 export function printPlace(place: Place): string {
   const items = [
+    place.abstractValue?.kind,
+    place.abstractValue ? ' ' : '',
     place.effect,
     ' ',
     printIdentifier(place.identifier),

compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts

@@ -268,6 +268,7 @@ function getManualMemoizationReplacement(
         kind: 'Identifier',
         identifier: fn.identifier,
         effect: Effect.Unknown,
+        abstractValue: null,
         reactive: false,
         loc,
       },
@@ -420,6 +421,7 @@ export function dropManualMemoization(func: HIRFunction): void {
                     kind: 'Identifier',
                     identifier: fnPlace.identifier,
                     effect: Effect.Unknown,
+                    abstractValue: null,
                     reactive: false,
                     loc: fnPlace.loc,
                   };

compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts

@@ -16,12 +16,14 @@ import {
   FunctionEffect,
   GeneratedSource,
   HIRFunction,
+  Identifier,
   IdentifierId,
   InstructionKind,
   InstructionValue,
   MethodCall,
   Phi,
   Place,
+  SourceLocation,
   SpreadPattern,
   Type,
   ValueKind,
@@ -45,6 +47,7 @@ import {
   eachTerminalOperand,
   eachTerminalSuccessor,
 } from '../HIR/visitors';
+import DisjointSet from '../Utils/DisjointSet';
 import {assertExhaustive} from '../Utils/utils';
 import {
   inferTerminalFunctionEffects,
@@ -103,7 +106,7 @@ const UndefinedValue: InstructionValue = {
 export default function inferReferenceEffects(
   fn: HIRFunction,
   options: {isFunctionExpression: boolean} = {isFunctionExpression: false},
-): void {
+): DisjointSet<IdentifierId> {
   /*
    * Initial state contains function params
    * TODO: include module declarations here as well
@@ -127,12 +130,14 @@ export default function inferReferenceEffects(
       properties: [],
       loc: ref.loc,
     };
-    initialState.initialize(value, {
+    const valueKind: AbstractValue = {
       kind: ValueKind.Context,
       reason: new Set([ValueReason.Other]),
       context: new Set([ref]),
-    });
+    };
+    initialState.initialize(value, valueKind);
     initialState.define(ref, value);
+    ref.abstractValue = valueKind;
   }
 
   const paramKind: AbstractValue = options.isFunctionExpression
@@ -177,12 +182,14 @@ export default function inferReferenceEffects(
           loc: ref.place.loc,
         };
       }
-      initialState.initialize(value, {
+      const valueKind: AbstractValue = {
         kind: ValueKind.Mutable,
         reason: new Set([ValueReason.Other]),
         context: new Set(),
-      });
+      };
+      initialState.initialize(value, valueKind);
       initialState.define(place, value);
+      place.abstractValue = valueKind;
     }
   } else {
     for (const param of fn.params) {
@@ -219,6 +226,7 @@ export default function inferReferenceEffects(
   }
   queue(fn.body.entry, initialState);
 
+  const finishedStates: Map<BlockId, InferenceState> = new Map();
   const functionEffects: Array<FunctionEffect> = fn.effects ?? [];
 
   while (queuedStates.size !== 0) {
@@ -231,6 +239,7 @@ export default function inferReferenceEffects(
 
       statesByBlock.set(blockId, incomingState);
       const state = incomingState.clone();
+      finishedStates.set(blockId, state);
       inferBlock(fn.env, state, block, functionEffects);
 
       for (const nextBlockId of eachTerminalSuccessor(block.terminal)) {
@@ -244,6 +253,12 @@ export default function inferReferenceEffects(
   } else {
     raiseFunctionEffectErrors(functionEffects);
   }
+
+  const summaryState = Array(...finishedStates.values()).reduce(
+    (acc, state) => acc.merge(state) ?? acc,
+  );
+
+  return summaryState.aliases;
 }
 
 type FreezeAction = {values: Set<InstructionValue>; reason: Set<ValueReason>};
@@ -261,18 +276,26 @@ class InferenceState {
    */
   #variables: Map<IdentifierId, Set<InstructionValue>>;
 
+  #aliases: DisjointSet<IdentifierId>;
+
   constructor(
     env: Environment,
     values: Map<InstructionValue, AbstractValue>,
     variables: Map<IdentifierId, Set<InstructionValue>>,
+    aliases: DisjointSet<IdentifierId>,
   ) {
     this.#env = env;
     this.#values = values;
     this.#variables = variables;
+    this.#aliases = aliases;
+  }
+
+  get aliases(): DisjointSet<IdentifierId> {
+    return this.#aliases;
   }
 
   static empty(env: Environment): InferenceState {
-    return new InferenceState(env, new Map(), new Map());
+    return new InferenceState(env, new Map(), new Map(), new DisjointSet());
   }
 
   // (Re)initializes a @param value with its default @param kind.
@@ -290,7 +313,7 @@ class InferenceState {
   values(place: Place): Array<InstructionValue> {
     const values = this.#variables.get(place.identifier.id);
     CompilerError.invariant(values != null, {
-      reason: `[hoisting] Expected value kind to be initialized`,
+      reason: `[hoisting] Expected value kind to be initialized in call to values()`,
       description: `${printPlace(place)}`,
       loc: place.loc,
       suggestions: null,
@@ -299,11 +322,11 @@ class InferenceState {
   }
 
   // Lookup the kind of the given @param value.
-  kind(place: Place): AbstractValue {
+  kind(place: {identifier: Identifier; loc: SourceLocation}): AbstractValue {
     const values = this.#variables.get(place.identifier.id);
     CompilerError.invariant(values != null, {
-      reason: `[hoisting] Expected value kind to be initialized`,
-      description: `${printPlace(place)}`,
+      reason: `[hoisting] Expected value kind to be initialized in call to kind()`,
+      description: `${place.identifier.id}`,
       loc: place.loc,
       suggestions: null,
     });
@@ -315,7 +338,7 @@ class InferenceState {
     }
     CompilerError.invariant(mergedKind !== null, {
       reason: `InferReferenceEffects::kind: Expected at least one value`,
-      description: `No value found at \`${printPlace(place)}\``,
+      description: `No value found at \`${place.identifier.id}\``,
       loc: place.loc,
       suggestions: null,
     });
@@ -332,6 +355,8 @@ class InferenceState {
       suggestions: null,
     });
     this.#variables.set(place.identifier.id, new Set(values));
+    this.#aliases.union([place.identifier.id, value.identifier.id]);
+    place.abstractValue = value.abstractValue;
   }
 
   // Defines (initializing or updating) a variable with a specific kind of value.
@@ -370,6 +395,7 @@ class InferenceState {
     reason: ValueReason,
   ): void {
     const values = this.#variables.get(place.identifier.id);
+    let valueKind: AbstractValue = this.kind(place);
     if (values === undefined) {
       CompilerError.invariant(effectKind !== Effect.Store, {
         reason: '[InferReferenceEffects] Unhandled store reference effect',
@@ -381,10 +407,12 @@ class InferenceState {
         effectKind === Effect.ConditionallyMutate
           ? Effect.ConditionallyMutate
           : Effect.Read;
+      place.abstractValue = valueKind;
       return;
     }
 
     const action = this.reference(place, effectKind, reason);
+    place.abstractValue = valueKind;
     action && freezeActions.push(action);
   }
 
@@ -431,7 +459,7 @@ class InferenceState {
       loc: place.loc,
       suggestions: null,
     });
-    let valueKind: AbstractValue | null = this.kind(place);
+    let valueKind: AbstractValue = this.kind(place);
     let effect: Effect | null = null;
     let freeze: null | FreezeAction = null;
     switch (effectKind) {
@@ -539,6 +567,7 @@ class InferenceState {
   merge(other: InferenceState): InferenceState | null {
     let nextValues: Map<InstructionValue, AbstractValue> | null = null;
     let nextVariables: Map<IdentifierId, Set<InstructionValue>> | null = null;
+    let nextAliases: DisjointSet<IdentifierId> | null = null;
 
     for (const [id, thisValue] of this.#values) {
       const otherValue = other.#values.get(id);
@@ -583,13 +612,21 @@ class InferenceState {
       nextVariables.set(id, new Set(otherValues));
     }
 
-    if (nextVariables === null && nextValues === null) {
+    if (!this.#aliases.equals(other.#aliases)) {
+      nextAliases = this.#aliases.copy();
+      for (const otherAliasSet of other.#aliases.buildSets()) {
+        nextAliases.union(Array(...otherAliasSet));
+      }
+    }
+
+    if (nextVariables === null && nextValues === null && nextAliases === null) {
       return null;
     } else {
       return new InferenceState(
         this.#env,
         nextValues ?? new Map(this.#values),
         nextVariables ?? new Map(this.#variables),
+        nextAliases ?? this.#aliases.copy(),
       );
     }
   }
@@ -604,6 +641,7 @@ class InferenceState {
       this.#env,
       new Map(this.#values),
       new Map(this.#variables),
+      this.#aliases.copy(),
     );
   }
 
@@ -634,10 +672,13 @@ class InferenceState {
 
   inferPhi(phi: Phi): void {
     const values: Set<InstructionValue> = new Set();
+    let valueKind;
     for (const [_, operand] of phi.operands) {
       const operandValues = this.#variables.get(operand.id);
       // This is a backedge that will be handled later by State.merge
       if (operandValues === undefined) continue;
+      const kind = this.kind({identifier: operand, loc: GeneratedSource});
+      valueKind = valueKind ? mergeAbstractValues(valueKind, kind) : kind;
       for (const v of operandValues) {
         values.add(v);
       }
@@ -645,6 +686,7 @@ class InferenceState {
 
     if (values.size > 0) {
       this.#variables.set(phi.id.id, values);
+      phi.abstractValue = valueKind!;
     }
   }
 }
@@ -673,6 +715,7 @@ function inferParam(
   }
   initialState.initialize(value, paramKind);
   initialState.define(place, value);
+  place.abstractValue = paramKind;
 }
 
 /*
@@ -913,6 +956,7 @@ function inferBlock(
 
         state.initialize(instrValue, valueKind);
         state.define(instr.lvalue, instrValue);
+        instr.lvalue.abstractValue = valueKind;
         instr.lvalue.effect = Effect.ConditionallyMutate;
         continuation = {kind: 'funeffects'};
         break;
@@ -972,6 +1016,7 @@ function inferBlock(
 
         state.initialize(instrValue, valueKind);
         state.define(instr.lvalue, instrValue);
+        instr.lvalue.abstractValue = valueKind;
         instr.lvalue.effect = Effect.Store;
         continuation = {kind: 'funeffects'};
         break;
@@ -1038,12 +1083,14 @@ function inferBlock(
           }
         }
 
-        state.initialize(instrValue, {
+        const valueKind: AbstractValue = {
           kind: ValueKind.Frozen,
           reason: new Set([ValueReason.Other]),
           context: new Set(),
-        });
+        };
+        state.initialize(instrValue, valueKind);
         state.define(instr.lvalue, instrValue);
+        instr.lvalue.abstractValue = valueKind;
         instr.lvalue.effect = Effect.ConditionallyMutate;
         continuation = {kind: 'funeffects'};
         break;
@@ -1156,12 +1203,14 @@ function inferBlock(
          * If a closure did not capture any mutable values, then we can consider it to be
          * frozen, which allows it to be independently memoized.
          */
-        state.initialize(instrValue, {
+        const valueKind: AbstractValue = {
           kind: hasMutableOperand ? ValueKind.Mutable : ValueKind.Frozen,
           reason: new Set([ValueReason.Other]),
           context: new Set(),
-        });
+        };
+        state.initialize(instrValue, valueKind);
         state.define(instr.lvalue, instrValue);
+        instr.lvalue.abstractValue = valueKind;
         instr.lvalue.effect = Effect.Store;
         continuation = {kind: 'funeffects'};
         break;
@@ -1204,6 +1253,7 @@ function inferBlock(
         );
         state.initialize(instrValue, returnValueKind);
         state.define(instr.lvalue, instrValue);
+        instr.lvalue.abstractValue = returnValueKind;
         instr.lvalue.effect = Effect.ConditionallyMutate;
         continuation = {kind: 'funeffects'};
         break;
@@ -1271,6 +1321,7 @@ function inferBlock(
 
         state.initialize(instrValue, returnValueKind);
         state.define(instr.lvalue, instrValue);
+        instr.lvalue.abstractValue = returnValueKind;
         instr.lvalue.effect = hasCaptureArgument
           ? Effect.Store
           : Effect.ConditionallyMutate;
@@ -1336,6 +1387,7 @@ function inferBlock(
           );
           state.initialize(instrValue, returnValueKind);
           state.define(instr.lvalue, instrValue);
+          instr.lvalue.abstractValue = returnValueKind;
           instr.lvalue.effect =
             instrValue.receiver.effect === Effect.Capture
               ? Effect.Store
@@ -1390,6 +1442,7 @@ function inferBlock(
 
         state.initialize(instrValue, returnValueKind);
         state.define(instr.lvalue, instrValue);
+        instr.lvalue.abstractValue = returnValueKind;
         instr.lvalue.effect = hasCaptureArgument
           ? Effect.Store
           : Effect.ConditionallyMutate;
@@ -1441,9 +1494,11 @@ function inferBlock(
           ValueReason.Other,
         );
         const lvalue = instr.lvalue;
+        const valueKind = state.kind(instrValue.object);
         lvalue.effect = Effect.ConditionallyMutate;
-        state.initialize(instrValue, state.kind(instrValue.object));
+        state.initialize(instrValue, valueKind);
         state.define(lvalue, instrValue);
+        instr.lvalue.abstractValue = valueKind;
         continuation = {kind: 'funeffects'};
         break;
       }
@@ -1490,12 +1545,14 @@ function inferBlock(
           Effect.Read,
           ValueReason.Other,
         );
-        state.initialize(instrValue, {
+        const valueKind: AbstractValue = {
           kind: ValueKind.Primitive,
           reason: new Set([ValueReason.Other]),
           context: new Set(),
-        });
+        };
+        state.initialize(instrValue, valueKind);
         state.define(instr.lvalue, instrValue);
+        instr.lvalue.abstractValue = valueKind;
         instr.lvalue.effect = Effect.Mutate;
         continuation = {kind: 'funeffects'};
         break;
@@ -1513,10 +1570,12 @@ function inferBlock(
           Effect.Read,
           ValueReason.Other,
         );
+        const valueKind = state.kind(instrValue.object);
         const lvalue = instr.lvalue;
         lvalue.effect = Effect.ConditionallyMutate;
-        state.initialize(instrValue, state.kind(instrValue.object));
+        state.initialize(instrValue, valueKind);
         state.define(lvalue, instrValue);
+        instr.lvalue.abstractValue = valueKind;
         continuation = {kind: 'funeffects'};
         break;
       }
@@ -1580,14 +1639,16 @@ function inferBlock(
             );
           }
         }
-        const lvalue = instr.lvalue;
-        lvalue.effect = Effect.ConditionallyMutate;
-        state.initialize(instrValue, {
+        const valueKind: AbstractValue = {
           kind: ValueKind.Frozen,
           reason: new Set([ValueReason.Other]),
           context: new Set(),
-        });
+        };
+        const lvalue = instr.lvalue;
+        lvalue.effect = Effect.ConditionallyMutate;
+        state.initialize(instrValue, valueKind);
         state.define(lvalue, instrValue);
+        instr.lvalue.abstractValue = valueKind;
         continuation = {kind: 'funeffects'};
         break;
       }
@@ -1622,13 +1683,13 @@ function inferBlock(
         const valueKind = state.kind(instrValue.place);
         state.initialize(instrValue, valueKind);
         state.define(lvalue, instrValue);
+        instr.lvalue.abstractValue = valueKind;
         continuation = {kind: 'funeffects'};
         break;
       }
       case 'DeclareLocal': {
         const value = UndefinedValue;
-        state.initialize(
-          value,
+        const valueKind: AbstractValue =
           // Catch params may be aliased to mutable values
           instrValue.lvalue.kind === InstructionKind.Catch
             ? {
@@ -1640,19 +1701,26 @@ function inferBlock(
                 kind: ValueKind.Primitive,
                 reason: new Set([ValueReason.Other]),
                 context: new Set(),
-              },
-        );
+              };
+        state.initialize(value, valueKind);
         state.define(instrValue.lvalue.place, value);
+        instrValue.lvalue.place.abstractValue = valueKind;
+        state.define(instr.lvalue, value);
+        instr.lvalue.abstractValue = valueKind;
         continuation = {kind: 'funeffects'};
         break;
       }
       case 'DeclareContext': {
-        state.initialize(instrValue, {
+        const valueKind: AbstractValue = {
           kind: ValueKind.Mutable,
           reason: new Set([ValueReason.Other]),
           context: new Set(),
-        });
+        };
+        state.initialize(instrValue, valueKind);
         state.define(instrValue.lvalue.place, instrValue);
+        instrValue.lvalue.place.abstractValue = valueKind;
+        state.define(instr.lvalue, instrValue);
+        instr.lvalue.abstractValue = valueKind;
         continuation = {kind: 'funeffects'};
         break;
       }
@@ -1740,6 +1808,10 @@ function inferBlock(
         );
         const lvalue = instr.lvalue;
         lvalue.effect = Effect.Store;
+        const valueKind = state.kind(instrValue.value);
+        state.initialize(instrValue, valueKind);
+        state.define(lvalue, instrValue);
+        instr.lvalue.abstractValue = valueKind;
         continuation = {kind: 'funeffects'};
         break;
       }
@@ -1853,8 +1925,10 @@ function inferBlock(
           Effect.Capture,
           ValueReason.Other,
         );
-        state.initialize(instrValue, state.kind(instrValue.collection));
+        const valueKind = state.kind(instrValue.collection);
+        state.initialize(instrValue, valueKind);
         state.define(instr.lvalue, instrValue);
+        instr.lvalue.abstractValue = valueKind;
         instr.lvalue.effect = Effect.Store;
         continuation = {kind: 'funeffects'};
         break;
@@ -1895,6 +1969,7 @@ function inferBlock(
 
       state.initialize(instrValue, continuation.valueKind);
       state.define(instr.lvalue, instrValue);
+      instr.lvalue.abstractValue = continuation.valueKind;
       instr.lvalue.effect = continuation.lvalueEffect ?? defaultLvalueEffect;
     }
 

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

@@ -237,6 +237,7 @@ class Transform extends ReactiveFunctionTransform<State> {
                   place: {
                     kind: 'Identifier',
                     effect: Effect.ConditionallyMutate,
+                    abstractValue: null,
                     loc,
                     reactive: true,
                     identifier: earlyReturnValue.value,
@@ -308,6 +309,7 @@ class Transform extends ReactiveFunctionTransform<State> {
                     kind: 'Identifier',
                     identifier: earlyReturnValue.value,
                     effect: Effect.Capture,
+                    abstractValue: null,
                     loc,
                     reactive: true,
                   },

compiler/packages/babel-plugin-react-compiler/src/SSA/EnterSSA.ts

@@ -191,6 +191,7 @@ class SSABuilder {
     const phi: Phi = {
       kind: 'Phi',
       id: newId,
+      abstractValue: null,
       operands: predDefs,
     };
 

compiler/packages/babel-plugin-react-compiler/src/Utils/DisjointSet.ts

@@ -124,6 +124,54 @@ export default class DisjointSet<T> {
     return [...sets.values()];
   }
 
+  copy(): DisjointSet<T> {
+    const copy = new DisjointSet<T>();
+    copy.#entries = new Map(this.#entries);
+    return copy;
+  }
+
+  equals(other: DisjointSet<T>): boolean {
+    if (this.size !== other.size) {
+      return false;
+    }
+    const rootMap = new Map<T, T>();
+    for (const thisGroupId of this.#entries.values()) {
+      const otherGroupId = other.find(thisGroupId);
+      if (otherGroupId === null || this.find(otherGroupId) !== thisGroupId) {
+        return false;
+      }
+      rootMap.set(thisGroupId, otherGroupId);
+    }
+
+    for (const otherGroupId of other.#entries.values()) {
+      if (!new Set(rootMap.values()).has(otherGroupId)) {
+        return false;
+      }
+    }
+
+    for (const item of this.#entries.keys()) {
+      const otherRoot = other.find(item);
+      if (otherRoot === null) {
+        return false;
+      }
+      const thisRoot = this.find(item);
+      CompilerError.invariant(thisRoot != null, {
+        reason: 'Expected item to be in set',
+        loc: null,
+      });
+      if (rootMap.get(thisRoot) !== otherRoot) {
+        return false;
+      }
+    }
+    for (const item of other.#entries.keys()) {
+      const thisRoot = this.find(item);
+      if (thisRoot === null) {
+        return false;
+      }
+    }
+    return true;
+  }
+
   get size(): number {
     return this.#entries.size;
   }

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

@@ -250,6 +250,7 @@ function validateInferredDep(
           identifier: dep.identifier,
           loc: GeneratedSource,
           effect: Effect.Read,
+          abstractValue: null,
           reactive: false,
         },
       },

compiler/packages/babel-plugin-react-compiler/src/__tests__/DisjointSet-test.ts

@@ -116,4 +116,45 @@ describe('DisjointSet', () => {
 
     identifiers.forEach((_, group) => expect(group).toBe(z));
   });
+
+  it('`.equals` is false when it should be', () => {
+    const x = new DisjointSet<TestIdentifier>();
+    const y = new DisjointSet<TestIdentifier>();
+
+    const [a, b] = makeIdentifiers('a', 'b');
+    x.union([a, b]);
+    y.union([a]);
+    y.union([b]);
+
+    expect(x.equals(y)).toBe(false);
+    expect(y.equals(x)).toBe(false);
+  });
+
+  it('`.equals` is true when it should be', () => {
+    const x = new DisjointSet<TestIdentifier>();
+    const y = new DisjointSet<TestIdentifier>();
+
+    const [a, b, c] = makeIdentifiers('a', 'b', 'c');
+    x.union([a, b, c]);
+    y.union([a, b]);
+    y.union([b, c]);
+
+    expect(x.equals(y)).toBe(true);
+    expect(y.equals(x)).toBe(true);
+  });
+
+  it('`.copy` doesnt mutate the underlying', () => {
+    const x = new DisjointSet<TestIdentifier>();
+
+    const [a, b] = makeIdentifiers('a', 'b');
+    x.union([a]);
+    x.union([b]);
+
+    const y = x.copy();
+
+    y.union([a, b]);
+
+    expect(x.find(a) !== x.find(b)).toBe(true);
+    expect(y.find(a) === y.find(b)).toBe(true);
+  });
 });

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.todo-repro-named-function-with-shadowed-local-same-name.expect.md

@@ -22,7 +22,7 @@ function Component(props) {
    7 |     return hasErrors;
    8 |   }
 >  9 |   return hasErrors();
-     |          ^^^^^^^^^ Invariant: [hoisting] Expected value for identifier to be initialized. hasErrors_0$16 (9:9)
+     |          ^^^^^^^^^ Invariant: [hoisting] Expected value kind to be initialized in call to kind(). 16 (9:9)
   10 | }
   11 |
 ```