[compiler] Propagate CreateFunction effects for functions that return functions

If you have a local helper function that itself returns a function (`() => () => { ... }`), we currently infer the return effect of the outer function as `Create mutable`. We correctly track the aliasing, but we lose some precision because we don't understand that a function specifically is being returned.

Here, we do some extra analysis of which values are returned in InferMutationAliasingRanges, and if the sole return value is a function we infer a `CreateFunction` effect. We also infer an `Assign` (instead of a Create) if the sole return value was one of the context variables or parameters.

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;
@@ -2541,10 +2498,47 @@ function computeEffectsForSignature(
         break;
       }
       case 'CreateFunction': {
-        CompilerError.throwTodo({
-          reason: `Support CreateFrom effects in signatures`,
-          loc: receiver.loc,
+        const applyInto = substitutions.get(effect.into.identifier.id);
+        if (applyInto == null || applyInto.length !== 1) {
+          return null;
+        }
+        const captures: Array<Place> = [];
+        for (let i = 0; i < effect.captures.length; i++) {
+          const substitution = substitutions.get(
+            effect.captures[i].identifier.id,
+          );
+          if (substitution == null || substitution.length !== 1) {
+            return null;
+          }
+          captures.push(substitution[0]);
+        }
+        const context: Array<Place> = [];
+        const originalContext = effect.function.loweredFunc.func.context;
+        for (let i = 0; i < originalContext.length; i++) {
+          const substitution = substitutions.get(
+            originalContext[i].identifier.id,
+          );
+          if (substitution == null || substitution.length !== 1) {
+            return null;
+          }
+          context.push(substitution[0]);
+        }
+        effects.push({
+          kind: 'CreateFunction',
+          into: applyInto[0],
+          function: {
+            ...effect.function,
+            loweredFunc: {
+              ...effect.function.loweredFunc,
+              func: {
+                ...effect.function.loweredFunc.func,
+                context,
+              },
+            },
+          },
+          captures,
         });
+        break;
       }
       default: {
         assertExhaustive(

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

@@ -141,7 +141,7 @@ export function inferMutationAliasingRanges(
         } else if (effect.kind === 'CreateFunction') {
           state.create(effect.into, {
             kind: 'Function',
-            function: effect.function.loweredFunc.func,
+            effect,
           });
         } else if (effect.kind === 'CreateFrom') {
           state.createFrom(index++, effect.from, effect.into);
@@ -156,7 +156,7 @@ export function inferMutationAliasingRanges(
            * invariant here.
            */
           if (!state.nodes.has(effect.into.identifier)) {
-            state.create(effect.into, {kind: 'Object'});
+            state.create(effect.into, {kind: 'Assign'});
           }
           state.assign(index++, effect.from, effect.into);
         } else if (effect.kind === 'Alias') {
@@ -474,6 +474,99 @@ export function inferMutationAliasingRanges(
     }
   }
 
+  const tracked: Array<Place> = [];
+  for (const param of [...fn.params, ...fn.context, fn.returns]) {
+    const place = param.kind === 'Identifier' ? param : param.place;
+    tracked.push(place);
+  }
+
+  const returned: Set<Node> = new Set();
+  const queue: Array<Node> = [state.nodes.get(fn.returns.identifier)!];
+  const seen: Set<Node> = new Set();
+  while (queue.length !== 0) {
+    const node = queue.pop()!;
+    if (seen.has(node)) {
+      continue;
+    }
+    seen.add(node);
+    for (const id of node.aliases.keys()) {
+      queue.push(state.nodes.get(id)!);
+    }
+    for (const id of node.createdFrom.keys()) {
+      queue.push(state.nodes.get(id)!);
+    }
+    if (node.id.id === fn.returns.identifier.id) {
+      continue;
+    }
+    switch (node.value.kind) {
+      case 'Assign':
+      case 'CreateFrom': {
+        break;
+      }
+      case 'Phi':
+      case 'Object':
+      case 'Function': {
+        returned.add(node);
+        break;
+      }
+      default: {
+        assertExhaustive(
+          node.value,
+          `Unexpected node value kind '${(node.value as any).kind}'`,
+        );
+      }
+    }
+  }
+  const returnedValues = [...returned];
+  if (
+    returnedValues.length === 1 &&
+    returnedValues[0].value.kind === 'Object' &&
+    tracked.some(place => place.identifier.id === returnedValues[0].id.id)
+  ) {
+    const from = tracked.find(
+      place => place.identifier.id === returnedValues[0].id.id,
+    )!;
+    functionEffects.push({
+      kind: 'Assign',
+      from,
+      into: fn.returns,
+    });
+  } else if (
+    returnedValues.length === 1 &&
+    returnedValues[0].value.kind === 'Function'
+  ) {
+    const outerContext = new Set(fn.context.map(p => p.identifier.id));
+    const effect = returnedValues[0].value.effect;
+    functionEffects.push({
+      kind: 'CreateFunction',
+      function: {
+        ...effect.function,
+        loweredFunc: {
+          func: {
+            ...effect.function.loweredFunc.func,
+            context: effect.function.loweredFunc.func.context.filter(p =>
+              outerContext.has(p.identifier.id),
+            ),
+          },
+        },
+      },
+      captures: effect.captures.filter(p => outerContext.has(p.identifier.id)),
+      into: fn.returns,
+    });
+  } else {
+    const returns = fn.returns.identifier;
+    functionEffects.push({
+      kind: 'Create',
+      into: fn.returns,
+      value: isPrimitiveType(returns)
+        ? ValueKind.Primitive
+        : isJsxType(returns.type)
+          ? ValueKind.Frozen
+          : ValueKind.Mutable,
+      reason: ValueReason.KnownReturnSignature,
+    });
+  }
+
   /**
    * Part 3
    * Finish populating the externally visible effects. Above we bubble-up the side effects
@@ -481,28 +574,12 @@ export function inferMutationAliasingRanges(
    * Here we populate an effect to create the return value as well as populating alias/capture
    * effects for how data flows between the params, context vars, and return.
    */
-  const returns = fn.returns.identifier;
-  functionEffects.push({
-    kind: 'Create',
-    into: fn.returns,
-    value: isPrimitiveType(returns)
-      ? ValueKind.Primitive
-      : isJsxType(returns.type)
-        ? ValueKind.Frozen
-        : ValueKind.Mutable,
-    reason: ValueReason.KnownReturnSignature,
-  });
   /**
    * Determine precise data-flow effects by simulating transitive mutations of the params/
    * captures and seeing what other params/context variables are affected. Anything that
    * would be transitively mutated needs a capture relationship.
    */
-  const tracked: Array<Place> = [];
   const ignoredErrors = new CompilerError();
-  for (const param of [...fn.params, ...fn.context, fn.returns]) {
-    const place = param.kind === 'Identifier' ? param : param.place;
-    tracked.push(place);
-  }
   for (const into of tracked) {
     const mutationIndex = index++;
     state.mutate(
@@ -588,9 +665,14 @@ type Node = {
   lastMutated: number;
   mutationReason: MutationReason | null;
   value:
+    | {kind: 'Assign'}
+    | {kind: 'CreateFrom'}
     | {kind: 'Object'}
     | {kind: 'Phi'}
-    | {kind: 'Function'; function: HIRFunction};
+    | {
+        kind: 'Function';
+        effect: Extract<AliasingEffect, {kind: 'CreateFunction'}>;
+      };
 };
 class AliasingState {
   nodes: Map<Identifier, Node> = new Map();
@@ -612,7 +694,7 @@ class AliasingState {
   }
 
   createFrom(index: number, from: Place, into: Place): void {
-    this.create(into, {kind: 'Object'});
+    this.create(into, {kind: 'CreateFrom'});
     const fromNode = this.nodes.get(from.identifier);
     const toNode = this.nodes.get(into.identifier);
     if (fromNode == null || toNode == null) {
@@ -674,7 +756,10 @@ class AliasingState {
         continue;
       }
       if (node.value.kind === 'Function') {
-        appendFunctionErrors(errors, node.value.function);
+        appendFunctionErrors(
+          errors,
+          node.value.effect.function.loweredFunc.func,
+        );
       }
       for (const [alias, when] of node.createdFrom) {
         if (when >= index) {
@@ -738,7 +823,10 @@ class AliasingState {
         node.transitive == null &&
         node.local == null
       ) {
-        appendFunctionErrors(errors, node.value.function);
+        appendFunctionErrors(
+          errors,
+          node.value.effect.function.loweredFunc.func,
+        );
       }
       if (transitive) {
         if (node.transitive == null || node.transitive.kind < kind) {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-returned-inner-fn-reassigns-context.expect.md

@@ -55,7 +55,7 @@ import { makeArray, Stringify, useIdentity } from "shared-runtime";
  */
 function Foo(t0) {
   "use memo";
-  const $ = _c(3);
+  const $ = _c(5);
   const { b } = t0;
 
   const fnFactory = () => () => {
@@ -66,18 +66,26 @@ function Foo(t0) {
   useIdentity();
 
   const fn = fnFactory();
-  const arr = makeArray(b);
-  fn(arr);
   let t1;
-  if ($[0] !== arr || $[1] !== myVar) {
-    t1 = <Stringify cb={myVar} value={arr} shouldInvokeFns={true} />;
-    $[0] = arr;
-    $[1] = myVar;
-    $[2] = t1;
+  if ($[0] !== b) {
+    t1 = makeArray(b);
+    $[0] = b;
+    $[1] = t1;
   } else {
-    t1 = $[2];
+    t1 = $[1];
   }
-  return t1;
+  const arr = t1;
+  fn(arr);
+  let t2;
+  if ($[2] !== arr || $[3] !== myVar) {
+    t2 = <Stringify cb={myVar} value={arr} shouldInvokeFns={true} />;
+    $[2] = arr;
+    $[3] = myVar;
+    $[4] = t2;
+  } else {
+    t2 = $[4];
+  }
+  return t2;
 }
 function _temp2() {
   return console.log("b");