[compiler] Effect inference across signatures and user-provided callbacks

This ties all the ideas together, showing the promise of the new inference. When applying effects, when we encounter an Apply effect we now check to see if the function we're calling is a locally defined FunctionExpression. If so, we construct a signature for it on the fly (we already have created the effects in AnalyzeFunctions), substitute in the args to get a set of effects we can apply, and then recursively apply those effects.

This required adding an ability for signatures to declare additional temporary places that they can reference. For example, Array.prototype.map needs a temporary to represent the items it extracts from the receiver array, and another temporary for the result of calling the user-provided function.

This also meant adding a `CreateFunction` effect which a) allows us to preserve the FunctionExpression value in the inference state (the previous Create effect meant we just created a dummy ObjectExpression) and b) allows dynamically constructing the ValueKind of the function based on whether it actually captures any mutable values.

Lots of other little fixes as well, such as changing function related effects (and PropertyLoad) to use Alias instead of Capture so that subsequent mutations of the output count as mutations of the input.

ghstack-source-id: 93c3bc3
Pull Request resolved: facebook/react#33384

Author: josephsavona

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

@@ -246,7 +246,7 @@ export const EnvironmentConfigSchema = z.object({
   /**
    * Enable a new model for mutability and aliasing inference
    */
-  enableNewMutationAliasingModel: z.boolean().default(false),
+  enableNewMutationAliasingModel: z.boolean().default(true),
 
   /**
    * Enables inference of optional dependency chains. Without this flag

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

@@ -192,6 +192,7 @@ export type FunctionSignature = {
   canonicalName?: string;
 
   aliasing?: AliasingSignature | null;
+  todo_aliasing?: AliasingSignature | null;
 };
 
 /*
@@ -320,6 +321,7 @@ addObject(BUILTIN_SHAPES, BuiltInArrayId, [
         params: [],
         rest: makeIdentifierId(1),
         returns: makeIdentifierId(2),
+        temporaries: [],
         effects: [
           // Push directly mutates the array itself
           {kind: 'Mutate', value: signatureArgument(0)},
@@ -367,7 +369,59 @@ addObject(BUILTIN_SHAPES, BuiltInArrayId, [
       returnValueKind: ValueKind.Mutable,
       noAlias: true,
       mutableOnlyIfOperandsAreMutable: true,
-      aliasing: null,
+      aliasing: {
+        receiver: makeIdentifierId(0),
+        params: [makeIdentifierId(1)],
+        rest: null,
+        returns: makeIdentifierId(2),
+        temporaries: [
+          // Temporary representing captured items of the receiver
+          signatureArgument(3),
+          // Temporary representing the result of the callback
+          signatureArgument(4),
+          /*
+           * Undefined `this` arg to the callback. Note the signature does not
+           * support passing an explicit thisArg second param
+           */
+          signatureArgument(5),
+        ],
+        effects: [
+          // Map creates a new mutable array
+          {
+            kind: 'Create',
+            into: signatureArgument(2),
+            value: ValueKind.Mutable,
+          },
+          // The first arg to the callback is an item extracted from the receiver array
+          {
+            kind: 'CreateFrom',
+            from: signatureArgument(0),
+            into: signatureArgument(3),
+          },
+          // The undefined this for the callback
+          {
+            kind: 'Create',
+            into: signatureArgument(5),
+            value: ValueKind.Primitive,
+          },
+          // calls the callback, returning the result into a temporary
+          {
+            kind: 'Apply',
+            receiver: signatureArgument(5),
+            args: [signatureArgument(3), {kind: 'Hole'}, signatureArgument(0)],
+            function: signatureArgument(1),
+            into: signatureArgument(4),
+            signature: null,
+            mutatesFunction: false,
+          },
+          // captures the result of the callback into the return array
+          {
+            kind: 'Capture',
+            from: signatureArgument(4),
+            into: signatureArgument(2),
+          },
+        ],
+      },
     }),
   ],
   [

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

@@ -972,6 +972,8 @@ export function printAliasingEffect(effect: AliasingEffect): string {
         .map(arg => {
           if (arg.kind === 'Identifier') {
             return printPlaceForAliasEffect(arg);
+          } else if (arg.kind === 'Hole') {
+            return ' ';
           }
           return `...${printPlaceForAliasEffect(arg.place)}`;
         })
@@ -986,6 +988,9 @@ export function printAliasingEffect(effect: AliasingEffect): string {
       }
       return `Apply ${printPlaceForAliasEffect(effect.into)} = ${receiverCallee}(${args})${signature != '' ? '\n     ' : ''}${signature}`;
     }
+    case 'CreateFunction': {
+      return `Function ${printPlaceForAliasEffect(effect.into)} = Function captures=[${effect.captures.map(printPlaceForAliasEffect).join(', ')}]`;
+    }
     case 'Freeze': {
       return `Freeze ${printPlaceForAliasEffect(effect.value)} ${effect.reason}`;
     }
@@ -1007,6 +1012,13 @@ function printPlaceForAliasEffect(place: Place): string {
 
 export function printAliasingSignature(signature: AliasingSignature): string {
   const tokens: Array<string> = ['function '];
+  if (signature.temporaries.length !== 0) {
+    tokens.push('<');
+    tokens.push(
+      signature.temporaries.map(temp => `$${temp.identifier.id}`).join(', '),
+    );
+    tokens.push('>');
+  }
   tokens.push('(');
   tokens.push('this=$' + String(signature.receiver));
   for (const param of signature.params) {

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

@@ -96,6 +96,7 @@ function lowerWithMutationAliasing(fn: HIRFunction): void {
         capturedOrMutated.add(effect.value.identifier.id);
         break;
       }
+      case 'CreateFunction':
       case 'Create':
       case 'Freeze':
       case 'ImmutableCapture': {