[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.

Author: josephsavona

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

@@ -2509,10 +2509,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

@@ -140,7 +140,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);
@@ -155,7 +155,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') {
@@ -469,35 +469,112 @@ 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
    * (MutateFrozen/MutableGlobal/Impure/Render) as well as mutations of context variables.
    * 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(
@@ -581,9 +658,14 @@ type Node = {
   local: {kind: MutationKind; loc: SourceLocation} | null;
   lastMutated: number;
   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();
@@ -604,7 +686,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) {
@@ -666,7 +748,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) {
@@ -728,7 +813,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 = $[1];
+  }
+  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 {
-    t1 = $[2];
+    t2 = $[4];
   }
-  return t1;
+  return t2;
 }
 function _temp2() {
   return console.log("b");