[compiler] moduleTypeProvider support for aliasing signatures

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

@@ -5,7 +5,14 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import {Effect, makeIdentifierId, ValueKind, ValueReason} from './HIR';
+import {
+  Effect,
+  GeneratedSource,
+  makeIdentifierId,
+  Place,
+  ValueKind,
+  ValueReason,
+} from './HIR';
 import {
   BUILTIN_SHAPES,
   BuiltInArrayId,
@@ -37,10 +44,15 @@ import {
   signatureArgument,
 } from './ObjectShape';
 import {BuiltInType, ObjectType, PolyType} from './Types';
-import {TypeConfig} from './TypeSchema';
+import {
+  AliasingEffectConfig,
+  AliasingSignatureConfig,
+  TypeConfig,
+} from './TypeSchema';
 import {assertExhaustive} from '../Utils/utils';
 import {isHookName} from './Environment';
 import {CompilerError, SourceLocation} from '..';
+import {AliasingEffect, AliasingSignature} from '../Inference/AliasingEffects';
 
 /*
  * This file exports types and defaults for JavaScript global objects.
@@ -891,6 +903,10 @@ export function installTypeConfig(
       }
     }
     case 'function': {
+      const aliasing =
+        typeConfig.aliasing != null
+          ? parseAliasingSignatureConfig(typeConfig.aliasing, moduleName, loc)
+          : null;
       return addFunction(shapes, [], {
         positionalParams: typeConfig.positionalParams,
         restParam: typeConfig.restParam,
@@ -906,9 +922,14 @@ export function installTypeConfig(
         noAlias: typeConfig.noAlias === true,
         mutableOnlyIfOperandsAreMutable:
           typeConfig.mutableOnlyIfOperandsAreMutable === true,
+        aliasing,
       });
     }
     case 'hook': {
+      const aliasing =
+        typeConfig.aliasing != null
+          ? parseAliasingSignatureConfig(typeConfig.aliasing, moduleName, loc)
+          : null;
       return addHook(shapes, {
         hookKind: 'Custom',
         positionalParams: typeConfig.positionalParams ?? [],
@@ -923,6 +944,7 @@ export function installTypeConfig(
         ),
         returnValueKind: typeConfig.returnValueKind ?? ValueKind.Frozen,
         noAlias: typeConfig.noAlias === true,
+        aliasing,
       });
     }
     case 'object': {
@@ -965,6 +987,90 @@ export function installTypeConfig(
   }
 }
 
+function parseAliasingSignatureConfig(
+  typeConfig: AliasingSignatureConfig,
+  moduleName: string,
+  loc: SourceLocation,
+): AliasingSignature {
+  const lifetimes = new Map<string, Place>();
+  function define(temp: string): Place {
+    CompilerError.invariant(!lifetimes.has(temp), {
+      reason: `Invalid type configuration for module`,
+      description: `Expected aliasing signature to have unique names for receiver, params, rest, returns, and temporaries in module '${moduleName}'`,
+      loc,
+    });
+    const place = signatureArgument(lifetimes.size);
+    lifetimes.set(temp, place);
+    return place;
+  }
+  function lookup(temp: string): Place {
+    const place = lifetimes.get(temp);
+    CompilerError.invariant(place != null, {
+      reason: `Invalid type configuration for module`,
+      description: `Expected aliasing signature effects to reference known names from receiver/params/rest/returns/temporaries, but '${temp}' is not a known name in '${moduleName}'`,
+      loc,
+    });
+    return place;
+  }
+  const receiver = define(typeConfig.receiver);
+  const params = typeConfig.params.map(define);
+  const rest = typeConfig.rest != null ? define(typeConfig.rest) : null;
+  const returns = define(typeConfig.returns);
+  const temporaries = typeConfig.temporaries.map(define);
+  const effects = typeConfig.effects.map(
+    (effect: AliasingEffectConfig): AliasingEffect => {
+      switch (effect.kind) {
+        case 'Assign': {
+          return {
+            kind: 'Assign',
+            from: lookup(effect.from),
+            into: lookup(effect.into),
+          };
+        }
+        case 'Create': {
+          return {
+            kind: 'Create',
+            into: lookup(effect.into),
+            reason: ValueReason.KnownReturnSignature,
+            value: effect.value,
+          };
+        }
+        case 'Freeze': {
+          return {
+            kind: 'Freeze',
+            value: lookup(effect.value),
+            reason: ValueReason.KnownReturnSignature,
+          };
+        }
+        case 'Impure': {
+          return {
+            kind: 'Impure',
+            place: lookup(effect.place),
+            error: CompilerError.throwTodo({
+              reason: 'Support impure effect declarations',
+              loc: GeneratedSource,
+            }),
+          };
+        }
+        default: {
+          assertExhaustive(
+            effect,
+            `Unexpected effect kind '${(effect as any).kind}'`,
+          );
+        }
+      }
+    },
+  );
+  return {
+    receiver: receiver.identifier.id,
+    params: params.map(p => p.identifier.id),
+    rest: rest != null ? rest.identifier.id : null,
+    returns: returns.identifier.id,
+    temporaries,
+    effects,
+  };
+}
+
 export function getReanimatedModuleType(registry: ShapeRegistry): ObjectType {
   // hooks that freeze args and return frozen value
   const frozenHooks = [

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

@@ -31,6 +31,86 @@ export const ObjectTypeSchema: z.ZodType<ObjectTypeConfig> = z.object({
   properties: ObjectPropertiesSchema.nullable(),
 });
 
+export const LifetimeIdSchema = z.string().refine(id => id.startsWith('@'), {
+  message: "Placeholder names must start with '@'",
+});
+
+export type FreezeEffectConfig = {
+  kind: 'Freeze';
+  value: string;
+};
+
+export const FreezeEffectSchema: z.ZodType<FreezeEffectConfig> = z.object({
+  kind: z.literal('Freeze'),
+  value: LifetimeIdSchema,
+});
+
+export type CreateEffectConfig = {
+  kind: 'Create';
+  into: string;
+  value: ValueKind;
+};
+
+export const CreateEffectSchema: z.ZodType<CreateEffectConfig> = z.object({
+  kind: z.literal('Create'),
+  into: LifetimeIdSchema,
+  value: ValueKindSchema,
+});
+
+export type AssignEffectConfig = {
+  kind: 'Assign';
+  from: string;
+  into: string;
+};
+
+export const AssignEffectSchema: z.ZodType<AssignEffectConfig> = z.object({
+  kind: z.literal('Assign'),
+  from: LifetimeIdSchema,
+  into: LifetimeIdSchema,
+});
+
+export type ImpureEffectConfig = {
+  kind: 'Impure';
+  place: string;
+};
+
+export const ImpureEffectSchema: z.ZodType<ImpureEffectConfig> = z.object({
+  kind: z.literal('Impure'),
+  place: LifetimeIdSchema,
+});
+
+export type AliasingEffectConfig =
+  | FreezeEffectConfig
+  | CreateEffectConfig
+  | AssignEffectConfig
+  | ImpureEffectConfig;
+
+export const AliasingEffectSchema: z.ZodType<AliasingEffectConfig> = z.union([
+  FreezeEffectSchema,
+  CreateEffectSchema,
+  AssignEffectSchema,
+  ImpureEffectSchema,
+]);
+
+export type AliasingSignatureConfig = {
+  receiver: string;
+  params: Array<string>;
+  rest: string | null;
+  returns: string;
+  effects: Array<AliasingEffectConfig>;
+  temporaries: Array<string>;
+};
+
+export const AliasingSignatureSchema: z.ZodType<AliasingSignatureConfig> =
+  z.object({
+    receiver: LifetimeIdSchema,
+    params: z.array(LifetimeIdSchema),
+    rest: LifetimeIdSchema.nullable(),
+    returns: LifetimeIdSchema,
+    effects: z.array(AliasingEffectSchema),
+    temporaries: z.array(LifetimeIdSchema),
+  });
+
 export type FunctionTypeConfig = {
   kind: 'function';
   positionalParams: Array<Effect>;
@@ -42,6 +122,7 @@ export type FunctionTypeConfig = {
   mutableOnlyIfOperandsAreMutable?: boolean | null | undefined;
   impure?: boolean | null | undefined;
   canonicalName?: string | null | undefined;
+  aliasing?: AliasingSignatureConfig | null | undefined;
 };
 export const FunctionTypeSchema: z.ZodType<FunctionTypeConfig> = z.object({
   kind: z.literal('function'),
@@ -54,6 +135,7 @@ export const FunctionTypeSchema: z.ZodType<FunctionTypeConfig> = z.object({
   mutableOnlyIfOperandsAreMutable: z.boolean().nullable().optional(),
   impure: z.boolean().nullable().optional(),
   canonicalName: z.string().nullable().optional(),
+  aliasing: AliasingSignatureSchema.nullable().optional(),
 });
 
 export type HookTypeConfig = {
@@ -63,6 +145,7 @@ export type HookTypeConfig = {
   returnType: TypeConfig;
   returnValueKind?: ValueKind | null | undefined;
   noAlias?: boolean | null | undefined;
+  aliasing?: AliasingSignatureConfig | null | undefined;
 };
 export const HookTypeSchema: z.ZodType<HookTypeConfig> = z.object({
   kind: z.literal('hook'),
@@ -71,6 +154,7 @@ export const HookTypeSchema: z.ZodType<HookTypeConfig> = z.object({
   returnType: z.lazy(() => TypeSchema),
   returnValueKind: ValueKindSchema.nullable().optional(),
   noAlias: z.boolean().nullable().optional(),
+  aliasing: AliasingSignatureSchema.nullable().optional(),
 });
 
 export type BuiltInTypeConfig =

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

@@ -8,6 +8,7 @@
 import {CompilerErrorDetailOptions} from '../CompilerError';
 import {
   FunctionExpression,
+  GeneratedSource,
   Hole,
   IdentifierId,
   ObjectMethod,
@@ -18,6 +19,7 @@ import {
   ValueReason,
 } from '../HIR';
 import {FunctionSignature} from '../HIR/ObjectShape';
+import {printSourceLocation} from '../HIR/PrintHIR';
 
 /**
  * `AliasingEffect` describes a set of "effects" that an instruction/terminal has on one or
@@ -200,10 +202,19 @@ export function hashEffect(effect: AliasingEffect): string {
       return [effect.kind, effect.value.identifier.id, effect.reason].join(':');
     }
     case 'Impure':
-    case 'Render':
+    case 'Render': {
+      return [effect.kind, effect.place.identifier.id].join(':');
+    }
     case 'MutateFrozen':
     case 'MutateGlobal': {
-      return [effect.kind, effect.place.identifier.id].join(':');
+      return [
+        effect.kind,
+        effect.place.identifier.id,
+        effect.error.severity,
+        effect.error.reason,
+        effect.error.description,
+        printSourceLocation(effect.error.loc ?? GeneratedSource),
+      ].join(':');
     }
     case 'Mutate':
     case 'MutateConditionally':