From 34179fe3449e141e980bbeaa8fc0a61b156113bb Mon Sep 17 00:00:00 2001 From: Joseph Savona <6425824+josephsavona@users.noreply.github.com> Date: Wed, 18 Jun 2025 15:43:48 -0700 Subject: [PATCH] [compiler] moduleTypeProvider support for aliasing signatures (#33526) This allows us to type things like `nullthrows()` or `identity()` functions where the return type is polymorphic on the input. --- [//]: # (BEGIN SAPLING FOOTER) Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/facebook/react/pull/33526). * #33571 * #33558 * #33547 * #33543 * #33533 * #33532 * #33530 * __->__ #33526 * #33522 * #33518 --- .../src/HIR/Globals.ts | 110 +++++++++++++++- .../src/HIR/TypeSchema.ts | 84 +++++++++++++ ...d-identity-function-frozen-input.expect.md | 117 ++++++++++++++++++ .../typed-identity-function-frozen-input.js | 39 ++++++ ...-identity-function-mutable-input.expect.md | 112 +++++++++++++++++ .../typed-identity-function-mutable-input.js | 35 ++++++ .../sprout/shared-runtime-type-provider.ts | 16 +++ .../snap/src/sprout/shared-runtime.ts | 4 + 8 files changed, 515 insertions(+), 2 deletions(-) create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-frozen-input.expect.md create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-frozen-input.js create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-mutable-input.expect.md create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-mutable-input.js diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/Globals.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/Globals.ts index c4c85be147..83f744cf68 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Globals.ts +++ b/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(); + 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 = [ diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/TypeSchema.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/TypeSchema.ts index 9aac2a264f..5ed39da0d9 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/TypeSchema.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/TypeSchema.ts @@ -31,6 +31,86 @@ export const ObjectTypeSchema: z.ZodType = 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 = z.object({ + kind: z.literal('Freeze'), + value: LifetimeIdSchema, +}); + +export type CreateEffectConfig = { + kind: 'Create'; + into: string; + value: ValueKind; +}; + +export const CreateEffectSchema: z.ZodType = z.object({ + kind: z.literal('Create'), + into: LifetimeIdSchema, + value: ValueKindSchema, +}); + +export type AssignEffectConfig = { + kind: 'Assign'; + from: string; + into: string; +}; + +export const AssignEffectSchema: z.ZodType = z.object({ + kind: z.literal('Assign'), + from: LifetimeIdSchema, + into: LifetimeIdSchema, +}); + +export type ImpureEffectConfig = { + kind: 'Impure'; + place: string; +}; + +export const ImpureEffectSchema: z.ZodType = z.object({ + kind: z.literal('Impure'), + place: LifetimeIdSchema, +}); + +export type AliasingEffectConfig = + | FreezeEffectConfig + | CreateEffectConfig + | AssignEffectConfig + | ImpureEffectConfig; + +export const AliasingEffectSchema: z.ZodType = z.union([ + FreezeEffectSchema, + CreateEffectSchema, + AssignEffectSchema, + ImpureEffectSchema, +]); + +export type AliasingSignatureConfig = { + receiver: string; + params: Array; + rest: string | null; + returns: string; + effects: Array; + temporaries: Array; +}; + +export const AliasingSignatureSchema: z.ZodType = + 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; @@ -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 = z.object({ kind: z.literal('function'), @@ -54,6 +135,7 @@ export const FunctionTypeSchema: z.ZodType = 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 = z.object({ kind: z.literal('hook'), @@ -71,6 +154,7 @@ export const HookTypeSchema: z.ZodType = z.object({ returnType: z.lazy(() => TypeSchema), returnValueKind: ValueKindSchema.nullable().optional(), noAlias: z.boolean().nullable().optional(), + aliasing: AliasingSignatureSchema.nullable().optional(), }); export type BuiltInTypeConfig = diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-frozen-input.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-frozen-input.expect.md new file mode 100644 index 0000000000..b15248df07 --- /dev/null +++ b/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 ; +} + +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 = ; + $[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)
{"inputs":[0],"output":{"a":0,"b":"value1","c":true}}
+
{"inputs":[1],"output":{"a":0,"b":"value1","c":true}}
+
{"inputs":[1],"output":{"a":0,"b":"value1","c":true}}
+
{"inputs":[0],"output":{"a":0,"b":"value1","c":true}}
+
{"inputs":[0],"output":{"a":0,"b":"value1","c":true}}
\ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-frozen-input.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-frozen-input.js new file mode 100644 index 0000000000..bcf6ecef02 --- /dev/null +++ b/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 ; +} + +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}, + ], +}; diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-mutable-input.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-mutable-input.expect.md new file mode 100644 index 0000000000..17fed05d93 --- /dev/null +++ b/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 ; +} + +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 = ; + $[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)
{"inputs":[0,0],"output":{"a":0,"b":"value1","c":true}}
+
{"inputs":[1,0],"output":{"a":0,"b":"value1","c":true}}
+
{"inputs":[1,1],"output":{"a":0,"b":"value1","c":true}}
+
{"inputs":[0,1],"output":{"a":0,"b":"value1","c":true}}
+
{"inputs":[0,0],"output":{"a":0,"b":"value1","c":true}}
\ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-mutable-input.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-mutable-input.js new file mode 100644 index 0000000000..719c89d11d --- /dev/null +++ b/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 ; +} + +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}, + ], +}; diff --git a/compiler/packages/snap/src/sprout/shared-runtime-type-provider.ts b/compiler/packages/snap/src/sprout/shared-runtime-type-provider.ts index 4c1d77f2f8..449bc8e688 100644 --- a/compiler/packages/snap/src/sprout/shared-runtime-type-provider.ts +++ b/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') { diff --git a/compiler/packages/snap/src/sprout/shared-runtime.ts b/compiler/packages/snap/src/sprout/shared-runtime.ts index 569d31cbd4..c8bf9272a2 100644 --- a/compiler/packages/snap/src/sprout/shared-runtime.ts +++ b/compiler/packages/snap/src/sprout/shared-runtime.ts @@ -396,4 +396,8 @@ export function typedLog(...values: Array): void { console.log(...values); } +export function typedIdentity(value: T): T { + return value; +} + export default typedLog;