[compiler] InferEffects uses effects as value keys

In InferReferenceEffects we used `InstructionValue` as the key to represent values, since each time we process an instruction this object will be the same. However this was always a bit of a hack, and in the new model and InferMutationAliasingEffects we can instead use the (creation) effect as the stable value. This avoids an extra layer of memoization since the effects are already interned anyway.

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

@@ -25,7 +25,6 @@ import {
   IdentifierId,
   Instruction,
   InstructionKind,
-  InstructionValue,
   isArrayType,
   isJsxType,
   isMapType,
@@ -57,7 +56,6 @@ import {
   printAliasingSignature,
   printIdentifier,
   printInstruction,
-  printInstructionValue,
   printPlace,
   printSourceLocation,
 } from '../HIR/PrintHIR';
@@ -108,11 +106,11 @@ export function inferMutationAliasingEffects(
   const statesByBlock: Map<BlockId, InferenceState> = new Map();
 
   for (const ref of fn.context) {
-    // TODO: using InstructionValue as a bit of a hack, but it's pragmatic
-    const value: InstructionValue = {
-      kind: 'ObjectExpression',
-      properties: [],
-      loc: ref.loc,
+    const value: AliasingEffect = {
+      kind: 'Create',
+      into: ref,
+      value: ValueKind.Context,
+      reason: ValueReason.Other,
     };
     initialState.initialize(value, {
       kind: ValueKind.Context,
@@ -145,10 +143,11 @@ export function inferMutationAliasingEffects(
     }
     if (ref != null) {
       const place = ref.kind === 'Identifier' ? ref : ref.place;
-      const value: InstructionValue = {
-        kind: 'ObjectExpression',
-        properties: [],
-        loc: place.loc,
+      const value: AliasingEffect = {
+        kind: 'Create',
+        into: place,
+        value: ValueKind.Mutable,
+        reason: ValueReason.Other,
       };
       initialState.initialize(value, {
         kind: ValueKind.Mutable,
@@ -265,8 +264,6 @@ function findHoistedContextDeclarations(
 class Context {
   internedEffects: Map<string, AliasingEffect> = new Map();
   instructionSignatureCache: Map<Instruction, InstructionSignature> = new Map();
-  effectInstructionValueCache: Map<AliasingEffect, InstructionValue> =
-    new Map();
   applySignatureCache: Map<
     AliasingSignature,
     Map<AliasingEffect, Array<AliasingEffect> | null>
@@ -318,10 +315,11 @@ function inferParam(
   paramKind: AbstractValue,
 ): void {
   const place = param.kind === 'Identifier' ? param : param.place;
-  const value: InstructionValue = {
-    kind: 'Primitive',
-    loc: place.loc,
-    value: undefined,
+  const value: AliasingEffect = {
+    kind: 'Create',
+    into: place,
+    value: paramKind.kind,
+    reason: ValueReason.Other,
   };
   initialState.initialize(value, paramKind);
   initialState.define(place, value);
@@ -542,20 +540,11 @@ function applyEffect(
       });
       initialized.add(effect.into.identifier.id);
 
-      let value = context.effectInstructionValueCache.get(effect);
-      if (value == null) {
-        value = {
-          kind: 'ObjectExpression',
-          properties: [],
-          loc: effect.into.loc,
-        };
-        context.effectInstructionValueCache.set(effect, value);
-      }
-      state.initialize(value, {
+      state.initialize(effect, {
         kind: effect.value,
         reason: new Set([effect.reason]),
       });
-      state.define(effect.into, value);
+      state.define(effect.into, effect);
       effects.push(effect);
       break;
     }
@@ -582,20 +571,11 @@ function applyEffect(
       initialized.add(effect.into.identifier.id);
 
       const fromValue = state.kind(effect.from);
-      let value = context.effectInstructionValueCache.get(effect);
-      if (value == null) {
-        value = {
-          kind: 'ObjectExpression',
-          properties: [],
-          loc: effect.into.loc,
-        };
-        context.effectInstructionValueCache.set(effect, value);
-      }
-      state.initialize(value, {
+      state.initialize(effect, {
         kind: fromValue.kind,
         reason: new Set(fromValue.reason),
       });
-      state.define(effect.into, value);
+      state.define(effect.into, effect);
       switch (fromValue.kind) {
         case ValueKind.Primitive:
         case ValueKind.Global: {
@@ -683,11 +663,11 @@ function applyEffect(
           operand.effect = Effect.Read;
         }
       }
-      state.initialize(effect.function, {
+      state.initialize(effect, {
         kind: isMutable ? ValueKind.Mutable : ValueKind.Frozen,
         reason: new Set([]),
       });
-      state.define(effect.into, effect.function);
+      state.define(effect.into, effect);
       for (const capture of effect.captures) {
         applyEffect(
           context,
@@ -793,38 +773,20 @@ function applyEffect(
             initialized,
             effects,
           );
-          let value = context.effectInstructionValueCache.get(effect);
-          if (value == null) {
-            value = {
-              kind: 'Primitive',
-              value: undefined,
-              loc: effect.from.loc,
-            };
-            context.effectInstructionValueCache.set(effect, value);
-          }
-          state.initialize(value, {
+          state.initialize(effect, {
             kind: fromKind,
             reason: new Set(fromValue.reason),
           });
-          state.define(effect.into, value);
+          state.define(effect.into, effect);
           break;
         }
         case ValueKind.Global:
         case ValueKind.Primitive: {
-          let value = context.effectInstructionValueCache.get(effect);
-          if (value == null) {
-            value = {
-              kind: 'Primitive',
-              value: undefined,
-              loc: effect.from.loc,
-            };
-            context.effectInstructionValueCache.set(effect, value);
-          }
-          state.initialize(value, {
+          state.initialize(effect, {
             kind: fromKind,
             reason: new Set(fromValue.reason),
           });
-          state.define(effect.into, value);
+          state.define(effect.into, effect);
           break;
         }
         default: {
@@ -839,14 +801,15 @@ function applyEffect(
       const functionValues = state.values(effect.function);
       if (
         functionValues.length === 1 &&
-        functionValues[0].kind === 'FunctionExpression' &&
-        functionValues[0].loweredFunc.func.aliasingEffects != null
+        functionValues[0].kind === 'CreateFunction' &&
+        functionValues[0].function.kind === 'FunctionExpression' &&
+        functionValues[0].function.loweredFunc.func.aliasingEffects != null
       ) {
         /*
          * We're calling a locally declared function, we already know it's effects!
          * We just have to substitute in the args for the params
          */
-        const functionExpr = functionValues[0];
+        const functionExpr = functionValues[0].function;
         let signature = context.functionSignatureCache.get(functionExpr);
         if (signature == null) {
           signature = buildSignatureFromFunctionExpression(
@@ -1136,19 +1099,19 @@ class InferenceState {
   #isFunctionExpression: boolean;
 
   // The kind of each value, based on its allocation site
-  #values: Map<InstructionValue, AbstractValue>;
+  #values: Map<AliasingEffect, AbstractValue>;
   /*
    * The set of values pointed to by each identifier. This is a set
    * to accomodate phi points (where a variable may have different
    * values from different control flow paths).
    */
-  #variables: Map<IdentifierId, Set<InstructionValue>>;
+  #variables: Map<IdentifierId, Set<AliasingEffect>>;
 
   constructor(
     env: Environment,
     isFunctionExpression: boolean,
-    values: Map<InstructionValue, AbstractValue>,
-    variables: Map<IdentifierId, Set<InstructionValue>>,
+    values: Map<AliasingEffect, AbstractValue>,
+    variables: Map<IdentifierId, Set<AliasingEffect>>,
   ) {
     this.env = env;
     this.#isFunctionExpression = isFunctionExpression;
@@ -1168,18 +1131,11 @@ class InferenceState {
   }
 
   // (Re)initializes a @param value with its default @param kind.
-  initialize(value: InstructionValue, kind: AbstractValue): void {
-    CompilerError.invariant(value.kind !== 'LoadLocal', {
-      reason:
-        '[InferMutationAliasingEffects] Expected all top-level identifiers to be defined as variables, not values',
-      description: null,
-      loc: value.loc,
-      suggestions: null,
-    });
+  initialize(value: AliasingEffect, kind: AbstractValue): void {
     this.#values.set(value, kind);
   }
 
-  values(place: Place): Array<InstructionValue> {
+  values(place: Place): Array<AliasingEffect> {
     const values = this.#variables.get(place.identifier.id);
     CompilerError.invariant(values != null, {
       reason: `[InferMutationAliasingEffects] Expected value kind to be initialized`,
@@ -1242,13 +1198,13 @@ class InferenceState {
   }
 
   // Defines (initializing or updating) a variable with a specific kind of value.
-  define(place: Place, value: InstructionValue): void {
+  define(place: Place, value: AliasingEffect): void {
     CompilerError.invariant(this.#values.has(value), {
       reason: `[InferMutationAliasingEffects] Expected value to be initialized at '${printSourceLocation(
-        value.loc,
+        place.loc,
       )}'`,
-      description: printInstructionValue(value),
-      loc: value.loc,
+      description: printAliasingEffect(value),
+      loc: place.loc,
       suggestions: null,
     });
     this.#variables.set(place.identifier.id, new Set([value]));
@@ -1288,17 +1244,17 @@ class InferenceState {
     }
   }
 
-  freezeValue(value: InstructionValue, reason: ValueReason): void {
+  freezeValue(value: AliasingEffect, reason: ValueReason): void {
     this.#values.set(value, {
       kind: ValueKind.Frozen,
       reason: new Set([reason]),
     });
     if (
-      value.kind === 'FunctionExpression' &&
+      value.kind === 'CreateFunction' &&
       (this.env.config.enablePreserveExistingMemoizationGuarantees ||
         this.env.config.enableTransitivelyFreezeFunctionExpressions)
     ) {
-      for (const place of value.loweredFunc.func.context) {
+      for (const place of value.function.loweredFunc.func.context) {
         this.freeze(place, reason);
       }
     }
@@ -1373,8 +1329,8 @@ class InferenceState {
    * termination.
    */
   merge(other: InferenceState): InferenceState | null {
-    let nextValues: Map<InstructionValue, AbstractValue> | null = null;
-    let nextVariables: Map<IdentifierId, Set<InstructionValue>> | null = null;
+    let nextValues: Map<AliasingEffect, AbstractValue> | null = null;
+    let nextVariables: Map<IdentifierId, Set<AliasingEffect>> | null = null;
 
     for (const [id, thisValue] of this.#values) {
       const otherValue = other.#values.get(id);
@@ -1398,7 +1354,7 @@ class InferenceState {
     for (const [id, thisValues] of this.#variables) {
       const otherValues = other.#variables.get(id);
       if (otherValues !== undefined) {
-        let mergedValues: Set<InstructionValue> | null = null;
+        let mergedValues: Set<AliasingEffect> | null = null;
         for (const otherValue of otherValues) {
           if (!thisValues.has(otherValue)) {
             mergedValues = mergedValues ?? new Set(thisValues);
@@ -1451,8 +1407,8 @@ class InferenceState {
    */
   debug(): any {
     const result: any = {values: {}, variables: {}};
-    const objects: Map<InstructionValue, number> = new Map();
-    function identify(value: InstructionValue): number {
+    const objects: Map<AliasingEffect, number> = new Map();
+    function identify(value: AliasingEffect): number {
       let id = objects.get(value);
       if (id == null) {
         id = objects.size;
@@ -1464,7 +1420,7 @@ class InferenceState {
       const id = identify(value);
       result.values[id] = {
         abstract: this.debugAbstractValue(kind),
-        value: printInstructionValue(value),
+        value: printAliasingEffect(value),
       };
     }
     for (const [variable, values] of this.#variables) {
@@ -1481,7 +1437,7 @@ class InferenceState {
   }
 
   inferPhi(phi: Phi): void {
-    const values: Set<InstructionValue> = new Set();
+    const values: Set<AliasingEffect> = new Set();
     for (const [_, operand] of phi.operands) {
       const operandValues = this.#variables.get(operand.identifier.id);
       // This is a backedge that will be handled later by State.merge
@@ -2347,8 +2303,9 @@ function areArgumentsImmutableAndNonMutating(
     const values = state.values(place);
     for (const value of values) {
       if (
-        value.kind === 'FunctionExpression' &&
-        value.loweredFunc.func.params.some(param => {
+        value.kind === 'CreateFunction' &&
+        value.function.kind === 'FunctionExpression' &&
+        value.function.loweredFunc.func.params.some(param => {
           const place = param.kind === 'Identifier' ? param : param.place;
           const range = place.identifier.mutableRange;
           return range.end > range.start + 1;