[compiler] moduleTypeProvider support for aliasing signatures

This allows us to type things like `nullthrows()` or `identity()` functions where the return type is polymorphic on the input.

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':

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

@@ -50,12 +50,14 @@ import {
 } from './InferReferenceEffects';
 import {
   assertExhaustive,
+  getOrInsertDefault,
   getOrInsertWith,
   Set_isSuperset,
 } from '../Utils/utils';
 import {
   printAliasingEffect,
   printAliasingSignature,
+  printFunction,
   printIdentifier,
   printInstruction,
   printInstructionValue,
@@ -195,12 +197,15 @@ export function inferMutationAliasingEffects(
   let count = 0;
   while (queuedStates.size !== 0) {
     count++;
-    if (count > 1000) {
+    if (count > 100) {
       console.log(
         'oops infinite loop',
         fn.id,
         typeof fn.loc !== 'symbol' ? fn.loc?.filename : null,
       );
+      if (DEBUG) {
+        console.log(printFunction(fn));
+      }
       throw new Error('infinite loop');
     }
     for (const [blockId, block] of fn.body.blocks) {
@@ -212,6 +217,11 @@ export function inferMutationAliasingEffects(
 
       statesByBlock.set(blockId, incomingState);
       const state = incomingState.clone();
+      if (DEBUG) {
+        console.log('*************');
+        console.log(`bb${block.id}`);
+        console.log('*************');
+      }
       inferBlock(context, state, block);
 
       for (const nextBlockId of eachTerminalSuccessor(block.terminal)) {
@@ -264,7 +274,13 @@ class Context {
   instructionSignatureCache: Map<Instruction, InstructionSignature> = new Map();
   effectInstructionValueCache: Map<AliasingEffect, InstructionValue> =
     new Map();
+  applySignatureCache: Map<
+    AliasingSignature,
+    Map<AliasingEffect, Array<AliasingEffect> | null>
+  > = new Map();
   catchHandlers: Map<BlockId, Place> = new Map();
+  functionSignatureCache: Map<FunctionExpression, AliasingSignature> =
+    new Map();
   isFuctionExpression: boolean;
   fn: HIRFunction;
   hoistedContextDeclarations: Map<DeclarationId, Place | null>;
@@ -279,6 +295,19 @@ class Context {
     this.hoistedContextDeclarations = hoistedContextDeclarations;
   }
 
+  cacheApplySignature(
+    signature: AliasingSignature,
+    effect: Extract<AliasingEffect, {kind: 'Apply'}>,
+    f: () => Array<AliasingEffect> | null,
+  ): Array<AliasingEffect> | null {
+    const inner = getOrInsertDefault(
+      this.applySignatureCache,
+      signature,
+      new Map(),
+    );
+    return getOrInsertWith(inner, effect, f);
+  }
+
   internEffect(effect: AliasingEffect): AliasingEffect {
     const hash = hashEffect(effect);
     let interned = this.internedEffects.get(hash);
@@ -352,11 +381,13 @@ function inferBlock(
           state.appendAlias(handlerParam, instr.lvalue);
           const kind = state.kind(instr.lvalue).kind;
           if (kind === ValueKind.Mutable || kind == ValueKind.Context) {
-            effects.push({
+            effects.push(
-              kind: 'Alias',
+              context.internEffect({
-              from: instr.lvalue,
+                kind: 'Alias',
-              into: handlerParam,
+                from: instr.lvalue,
-            });
+                into: handlerParam,
+              }),
+            );
           }
         }
       }
@@ -365,11 +396,11 @@ function inferBlock(
   } else if (terminal.kind === 'return') {
     if (!context.isFuctionExpression) {
       terminal.effects = [
-        {
+        context.internEffect({
           kind: 'Freeze',
           value: terminal.value,
           reason: ValueReason.JsxCaptured,
-        },
+        }),
       ];
     }
   }
@@ -546,20 +577,21 @@ function applyEffect(
           break;
         }
         case ValueKind.Frozen: {
-          effects.push({
+          applyEffect(
-            kind: 'ImmutableCapture',
+            context,
-            from: effect.from,
+            state,
-            into: effect.into,
+            {
-          });
+              kind: 'ImmutableCapture',
+              from: effect.from,
+              into: effect.into,
+            },
+            aliased,
+            effects,
+          );
           break;
         }
         default: {
-          effects.push({
+          effects.push(effect);
-            // OK: recording information flow
-            kind: 'CreateFrom', // prev Alias
-            from: effect.from,
-            into: effect.into,
-          });
         }
       }
       break;
@@ -658,11 +690,17 @@ function applyEffect(
         }
         case ValueKind.Frozen: {
           isMutableReferenceType = false;
-          effects.push({
+          applyEffect(
-            kind: 'ImmutableCapture',
+            context,
-            from: effect.from,
+            state,
-            into: effect.into,
+            {
-          });
+              kind: 'ImmutableCapture',
+              from: effect.from,
+              into: effect.into,
+            },
+            aliased,
+            effects,
+          );
           break;
         }
         default: {
@@ -684,11 +722,17 @@ function applyEffect(
       const fromKind = fromValue.kind;
       switch (fromKind) {
         case ValueKind.Frozen: {
-          effects.push({
+          applyEffect(
-            kind: 'ImmutableCapture',
+            context,
-            from: effect.from,
+            state,
-            into: effect.into,
+            {
-          });
+              kind: 'ImmutableCapture',
+              from: effect.from,
+              into: effect.into,
+            },
+            aliased,
+            effects,
+          );
           let value = context.effectInstructionValueCache.get(effect);
           if (value == null) {
             value = {
@@ -746,23 +790,33 @@ function applyEffect(
          * 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 signature = buildSignatureFromFunctionExpression(
+        const functionExpr = functionValues[0];
-          state.env,
+        let signature = context.functionSignatureCache.get(functionExpr);
-          functionValues[0],
+        if (signature == null) {
-        );
+          signature = buildSignatureFromFunctionExpression(
+            state.env,
+            functionExpr,
+          );
+          context.functionSignatureCache.set(functionExpr, signature);
+        }
         if (DEBUG) {
           console.log(
             `constructed alias signature:\n${printAliasingSignature(signature)}`,
           );
         }
-        const signatureEffects = computeEffectsForSignature(
+        const signatureEffects = context.cacheApplySignature(
-          state.env,
           signature,
-          effect.into,
+          effect,
-          effect.receiver,
+          () =>
-          effect.args,
+            computeEffectsForSignature(
-          functionValues[0].loweredFunc.func.context,
+              state.env,
-          effect.loc,
+              signature,
+              effect.into,
+              effect.receiver,
+              effect.args,
+              functionExpr.loweredFunc.func.context,
+              effect.loc,
+            ),
         );
         if (signatureEffects != null) {
           if (DEBUG) {
@@ -781,18 +835,24 @@ function applyEffect(
           break;
         }
       }
-      const signatureEffects =
+      let signatureEffects = null;
-        effect.signature?.aliasing != null
+      if (effect.signature?.aliasing != null) {
-          ? computeEffectsForSignature(
+        const signature = effect.signature.aliasing;
+        signatureEffects = context.cacheApplySignature(
+          effect.signature.aliasing,
+          effect,
+          () =>
+            computeEffectsForSignature(
               state.env,
-              effect.signature.aliasing,
+              signature,
               effect.into,
               effect.receiver,
               effect.args,
               [],
               effect.loc,
-            )
+            ),
-          : null;
+        );
+      }
       if (signatureEffects != null) {
         if (DEBUG) {
           console.log('apply aliasing signature effects');
@@ -935,30 +995,42 @@ function applyEffect(
             effect.value.identifier.declarationId,
           );
           if (hoistedAccess != null && hoistedAccess.loc != effect.value.loc) {
-            effects.push({
+            applyEffect(
+              context,
+              state,
+              {
+                kind: 'MutateFrozen',
+                place: effect.value,
+                error: {
+                  severity: ErrorSeverity.InvalidReact,
+                  reason: `This variable is accessed before it is declared, which may prevent it from updating as the assigned value changes over time`,
+                  description,
+                  loc: hoistedAccess.loc,
+                  suggestions: null,
+                },
+              },
+              aliased,
+              effects,
+            );
+          }
+
+          applyEffect(
+            context,
+            state,
+            {
               kind: 'MutateFrozen',
               place: effect.value,
               error: {
                 severity: ErrorSeverity.InvalidReact,
-                reason: `This variable is accessed before it is declared, which may prevent it from updating as the assigned value changes over time`,
+                reason: `This variable is accessed before it is declared, which prevents the earlier access from updating when this value changes over time`,
                 description,
-                loc: hoistedAccess.loc,
+                loc: effect.value.loc,
                 suggestions: null,
               },
-            });
-          }
-
-          effects.push({
-            kind: 'MutateFrozen',
-            place: effect.value,
-            error: {
-              severity: ErrorSeverity.InvalidReact,
-              reason: `This variable is accessed before it is declared, which prevents the earlier access from updating when this value changes over time`,
-              description,
-              loc: effect.value.loc,
-              suggestions: null,
             },
-          });
+            aliased,
+            effects,
+          );
         } else {
           const reason = getWriteErrorReason({
             kind: value.kind,
@@ -970,18 +1042,26 @@ function applyEffect(
             effect.value.identifier.name.kind === 'named'
               ? `Found mutation of \`${effect.value.identifier.name.value}\``
               : null;
-          effects.push({
+          applyEffect(
-            kind:
+            context,
-              value.kind === ValueKind.Frozen ? 'MutateFrozen' : 'MutateGlobal',
+            state,
-            place: effect.value,
+            {
-            error: {
+              kind:
-              severity: ErrorSeverity.InvalidReact,
+                value.kind === ValueKind.Frozen
-              reason,
+                  ? 'MutateFrozen'
-              description,
+                  : 'MutateGlobal',
-              loc: effect.value.loc,
+              place: effect.value,
-              suggestions: null,
+              error: {
+                severity: ErrorSeverity.InvalidReact,
+                reason,
+                description,
+                loc: effect.value.loc,
+                suggestions: null,
+              },
             },
-          });
+            aliased,
+            effects,
+          );
         }
       }
       break;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/repro-compiler-infinite-loop.expect.md

@@ -0,0 +1,54 @@
+
+## Input
+
+```javascript
+// @flow @enableNewMutationAliasingModel
+
+import fbt from 'fbt';
+
+component Component() {
+  const sections = Object.keys(items);
+
+  for (let i = 0; i < sections.length; i += 3) {
+    chunks.push(
+      sections.slice(i, i + 3).map(section => {
+        return <Child />;
+      })
+    );
+  }
+
+  return <Child />;
+}
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+
+import fbt from "fbt";
+
+function Component() {
+  const $ = _c(1);
+  const sections = Object.keys(items);
+  for (let i = 0; i < sections.length; i = i + 3, i) {
+    chunks.push(sections.slice(i, i + 3).map(_temp));
+  }
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = <Child />;
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+function _temp(section) {
+  return <Child />;
+}
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/repro-compiler-infinite-loop.js

@@ -0,0 +1,17 @@
+// @flow @enableNewMutationAliasingModel
+
+import fbt from 'fbt';
+
+component Component() {
+  const sections = Object.keys(items);
+
+  for (let i = 0; i < sections.length; i += 3) {
+    chunks.push(
+      sections.slice(i, i + 3).map(section => {
+        return <Child />;
+      })
+    );
+  }
+
+  return <Child />;
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/repro-jsx-captures-value-mutated-later.expect.md

@@ -0,0 +1,53 @@
+
+## Input
+
+```javascript
+// @flow @enableNewMutationAliasingModel
+
+import {identity, Stringify, useFragment} from 'shared-runtime';
+
+component Example() {
+  const data = useFragment();
+
+  const {a, b} = identity(data);
+
+  const el = <Stringify tooltip={b} />;
+
+  identity(a.at(0));
+
+  return <Stringify icon={el} />;
+}
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+
+import { identity, Stringify, useFragment } from "shared-runtime";
+
+function Example() {
+  const $ = _c(2);
+  const data = useFragment();
+  let t0;
+  if ($[0] !== data) {
+    const { a, b } = identity(data);
+
+    const el = <Stringify tooltip={b} />;
+
+    identity(a.at(0));
+
+    t0 = <Stringify icon={el} />;
+    $[0] = data;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
+}
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/repro-jsx-captures-value-mutated-later.js

@@ -0,0 +1,15 @@
+// @flow @enableNewMutationAliasingModel
+
+import {identity, Stringify, useFragment} from 'shared-runtime';
+
+component Example() {
+  const data = useFragment();
+
+  const {a, b} = identity(data);
+
+  const el = <Stringify tooltip={b} />;
+
+  identity(a.at(0));
+
+  return <Stringify icon={el} />;
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-frozen-input.expect.md

@@ -0,0 +1,117 @@
+
+## Input
+
+```javascript
+// @enableNewMutationAliasingModel
+
+import {
+  identity,
+  makeObject_Primitives,
+  typedIdentity,
+  useIdentity,
+  ValidateMemoization,
+} from 'shared-runtime';
+
+function Component({a, b}) {
+  // create a mutable value with input `a`
+  const x = makeObject_Primitives(a);
+
+  // freeze the value
+  useIdentity(x);
+
+  // known to pass-through via aliasing signature
+  const x2 = typedIdentity(x);
+
+  // Unknown function so we assume it conditionally mutates,
+  // but x2 is frozen so this downgrades to a read.
+  // x should *not* take b as a dependency
+  identity(x2, b);
+
+  return <ValidateMemoization inputs={[a]} output={x} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 0, b: 0}],
+  sequentialRenders: [
+    {a: 0, b: 0},
+    {a: 1, b: 0},
+    {a: 1, b: 1},
+    {a: 0, b: 1},
+    {a: 0, b: 0},
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enableNewMutationAliasingModel
+
+import {
+  identity,
+  makeObject_Primitives,
+  typedIdentity,
+  useIdentity,
+  ValidateMemoization,
+} from "shared-runtime";
+
+function Component(t0) {
+  const $ = _c(7);
+  const { a, b } = t0;
+  let t1;
+  if ($[0] !== a) {
+    t1 = makeObject_Primitives(a);
+    $[0] = a;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  const x = t1;
+
+  useIdentity(x);
+
+  const x2 = typedIdentity(x);
+
+  identity(x2, b);
+  let t2;
+  if ($[2] !== a) {
+    t2 = [a];
+    $[2] = a;
+    $[3] = t2;
+  } else {
+    t2 = $[3];
+  }
+  let t3;
+  if ($[4] !== t2 || $[5] !== x) {
+    t3 = <ValidateMemoization inputs={t2} output={x} />;
+    $[4] = t2;
+    $[5] = x;
+    $[6] = t3;
+  } else {
+    t3 = $[6];
+  }
+  return t3;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ a: 0, b: 0 }],
+  sequentialRenders: [
+    { a: 0, b: 0 },
+    { a: 1, b: 0 },
+    { a: 1, b: 1 },
+    { a: 0, b: 1 },
+    { a: 0, b: 0 },
+  ],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"inputs":[0],"output":{"a":0,"b":"value1","c":true}}</div>
+<div>{"inputs":[1],"output":{"a":0,"b":"value1","c":true}}</div>
+<div>{"inputs":[1],"output":{"a":0,"b":"value1","c":true}}</div>
+<div>{"inputs":[0],"output":{"a":0,"b":"value1","c":true}}</div>
+<div>{"inputs":[0],"output":{"a":0,"b":"value1","c":true}}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-frozen-input.js

@@ -0,0 +1,39 @@
+// @enableNewMutationAliasingModel
+
+import {
+  identity,
+  makeObject_Primitives,
+  typedIdentity,
+  useIdentity,
+  ValidateMemoization,
+} from 'shared-runtime';
+
+function Component({a, b}) {
+  // create a mutable value with input `a`
+  const x = makeObject_Primitives(a);
+
+  // freeze the value
+  useIdentity(x);
+
+  // known to pass-through via aliasing signature
+  const x2 = typedIdentity(x);
+
+  // Unknown function so we assume it conditionally mutates,
+  // but x2 is frozen so this downgrades to a read.
+  // x should *not* take b as a dependency
+  identity(x2, b);
+
+  return <ValidateMemoization inputs={[a]} output={x} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 0, b: 0}],
+  sequentialRenders: [
+    {a: 0, b: 0},
+    {a: 1, b: 0},
+    {a: 1, b: 1},
+    {a: 0, b: 1},
+    {a: 0, b: 0},
+  ],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-mutable-input.expect.md

@@ -0,0 +1,112 @@
+
+## Input
+
+```javascript
+// @enableNewMutationAliasingModel
+
+import {
+  identity,
+  makeObject_Primitives,
+  typedIdentity,
+  useIdentity,
+  ValidateMemoization,
+} from 'shared-runtime';
+
+function Component({a, b}) {
+  // create a mutable value with input `a`
+  const x = makeObject_Primitives(a);
+
+  // known to pass-through via aliasing signature
+  const x2 = typedIdentity(x);
+
+  // Unknown function so we assume it conditionally mutates,
+  // and x is still mutable so
+  identity(x2, b);
+
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 0, b: 0}],
+  sequentialRenders: [
+    {a: 0, b: 0},
+    {a: 1, b: 0},
+    {a: 1, b: 1},
+    {a: 0, b: 1},
+    {a: 0, b: 0},
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enableNewMutationAliasingModel
+
+import {
+  identity,
+  makeObject_Primitives,
+  typedIdentity,
+  useIdentity,
+  ValidateMemoization,
+} from "shared-runtime";
+
+function Component(t0) {
+  const $ = _c(9);
+  const { a, b } = t0;
+  let x;
+  if ($[0] !== a || $[1] !== b) {
+    x = makeObject_Primitives(a);
+
+    const x2 = typedIdentity(x);
+
+    identity(x2, b);
+    $[0] = a;
+    $[1] = b;
+    $[2] = x;
+  } else {
+    x = $[2];
+  }
+  let t1;
+  if ($[3] !== a || $[4] !== b) {
+    t1 = [a, b];
+    $[3] = a;
+    $[4] = b;
+    $[5] = t1;
+  } else {
+    t1 = $[5];
+  }
+  let t2;
+  if ($[6] !== t1 || $[7] !== x) {
+    t2 = <ValidateMemoization inputs={t1} output={x} />;
+    $[6] = t1;
+    $[7] = x;
+    $[8] = t2;
+  } else {
+    t2 = $[8];
+  }
+  return t2;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ a: 0, b: 0 }],
+  sequentialRenders: [
+    { a: 0, b: 0 },
+    { a: 1, b: 0 },
+    { a: 1, b: 1 },
+    { a: 0, b: 1 },
+    { a: 0, b: 0 },
+  ],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"inputs":[0,0],"output":{"a":0,"b":"value1","c":true}}</div>
+<div>{"inputs":[1,0],"output":{"a":0,"b":"value1","c":true}}</div>
+<div>{"inputs":[1,1],"output":{"a":0,"b":"value1","c":true}}</div>
+<div>{"inputs":[0,1],"output":{"a":0,"b":"value1","c":true}}</div>
+<div>{"inputs":[0,0],"output":{"a":0,"b":"value1","c":true}}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-mutable-input.js

@@ -0,0 +1,35 @@
+// @enableNewMutationAliasingModel
+
+import {
+  identity,
+  makeObject_Primitives,
+  typedIdentity,
+  useIdentity,
+  ValidateMemoization,
+} from 'shared-runtime';
+
+function Component({a, b}) {
+  // create a mutable value with input `a`
+  const x = makeObject_Primitives(a);
+
+  // known to pass-through via aliasing signature
+  const x2 = typedIdentity(x);
+
+  // Unknown function so we assume it conditionally mutates,
+  // and x is still mutable so
+  identity(x2, b);
+
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 0, b: 0}],
+  sequentialRenders: [
+    {a: 0, b: 0},
+    {a: 1, b: 0},
+    {a: 1, b: 1},
+    {a: 0, b: 1},
+    {a: 0, b: 0},
+  ],
+};

compiler/packages/snap/src/sprout/shared-runtime-type-provider.ts

@@ -69,6 +69,22 @@ export function makeSharedRuntimeTypeProvider({
             returnValueKind: ValueKindEnum.Mutable,
             noAlias: true,
           },
+          typedIdentity: {
+            kind: 'function',
+            positionalParams: [EffectEnum.Read],
+            restParam: null,
+            calleeEffect: EffectEnum.Read,
+            returnType: {kind: 'type', name: 'Any'},
+            returnValueKind: ValueKindEnum.Mutable,
+            aliasing: {
+              receiver: '@receiver',
+              params: ['@value'],
+              rest: null,
+              returns: '@return',
+              temporaries: [],
+              effects: [{kind: 'Assign', from: '@value', into: '@return'}],
+            },
+          },
         },
       };
     } else if (moduleName === 'ReactCompilerTest') {

compiler/packages/snap/src/sprout/shared-runtime.ts

@@ -396,4 +396,8 @@ export function typedLog(...values: Array<any>): void {
   console.log(...values);
 }
 
+export function typedIdentity<T>(value: T): T {
+  return value;
+}
+
 export default typedLog;