diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts index 9a59397126cc5..277921d5e25ee 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts @@ -407,6 +407,23 @@ const EnvironmentConfigSchema = z.object({ * and identifiers have been changed. */ hookPattern: z.string().nullable().default(null), + + /** + * If enabled, this will treat objects named as `ref` or if their names end with the substring `Ref`, + * and contain a property named `current`, as React refs. + * + * ``` + * const ref = useMyRef(); + * const myRef = useMyRef2(); + * useEffect(() => { + * ref.current = ...; + * myRef.current = ...; + * }) + * ``` + * + * Here the variables `ref` and `myRef` will be typed as Refs. + */ + enableTreatRefLikeIdentifiersAsRefs: z.boolean().nullable().default(false), }); export type EnvironmentConfig = z.infer; diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/Types.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/Types.ts index f1d08f762b9a9..be9ce00013c14 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Types.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/Types.ts @@ -57,7 +57,8 @@ export type PhiType = { }; export type PropType = { kind: "Property"; - object: Type; + objectType: Type; + objectName: string; propertyName: string; }; @@ -124,7 +125,8 @@ export function duplicateType(type: Type): Type { case "Property": { return { kind: "Property", - object: duplicateType(type.object), + objectType: duplicateType(type.objectType), + objectName: type.objectName, propertyName: type.propertyName, }; } @@ -165,11 +167,13 @@ function objectMethodTypeEquals(tA: Type, tB: Type): boolean { function propTypeEquals(tA: Type, tB: Type): boolean { if (tA.kind === "Property" && tB.kind === "Property") { - if (!typeEquals(tA.object, tB.object)) { + if (!typeEquals(tA.objectType, tB.objectType)) { return false; } - return tA.propertyName === tB.propertyName; + return ( + tA.propertyName === tB.propertyName && tA.objectName === tB.objectName + ); } return false; diff --git a/compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts b/compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts index b77d7b861c6ea..89bbde944f078 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts @@ -11,8 +11,11 @@ import { Environment } from "../HIR"; import { lowerType } from "../HIR/BuildHIR"; import { HIRFunction, + Identifier, + IdentifierId, Instruction, makeType, + PropType, Type, typeEquals, TypeId, @@ -24,6 +27,7 @@ import { BuiltInJsxId, BuiltInObjectId, BuiltInPropsId, + BuiltInRefValueId, BuiltInUseRefId, } from "../HIR/ObjectShape"; import { eachInstructionLValue, eachInstructionOperand } from "../HIR/visitors"; @@ -117,6 +121,7 @@ function* generate( } } + const names = new Map(); for (const [_, block] of func.body.blocks) { for (const phi of block.phis) { yield equation(phi.type, { @@ -126,13 +131,28 @@ function* generate( } for (const instr of block.instructions) { - yield* generateInstructionTypes(func.env, instr); + yield* generateInstructionTypes(func.env, names, instr); } } } +function setName( + names: Map, + id: IdentifierId, + name: Identifier +): void { + if (name.name?.kind === "named") { + names.set(id, name.name.value); + } +} + +function getName(names: Map, id: IdentifierId): string { + return names.get(id) ?? ""; +} + function* generateInstructionTypes( env: Environment, + names: Map, instr: Instruction ): Generator { const { lvalue, value } = instr; @@ -152,6 +172,7 @@ function* generateInstructionTypes( } case "LoadLocal": { + setName(names, lvalue.identifier.id, value.place.identifier); yield equation(left, value.place.identifier.type); break; } @@ -250,7 +271,8 @@ function* generateInstructionTypes( case "PropertyLoad": { yield equation(left, { kind: "Property", - object: value.object.identifier.type, + objectType: value.object.identifier.type, + objectName: getName(names, value.object.identifier.id), propertyName: value.property, }); break; @@ -278,7 +300,8 @@ function* generateInstructionTypes( const propertyName = String(i); yield equation(item.identifier.type, { kind: "Property", - object: value.value.identifier.type, + objectType: value.value.identifier.type, + objectName: getName(names, value.value.identifier.id), propertyName, }); } else { @@ -294,7 +317,8 @@ function* generateInstructionTypes( ) { yield equation(property.place.identifier.type, { kind: "Property", - object: value.value.identifier.type, + objectType: value.value.identifier.type, + objectName: getName(names, value.value.identifier.id), propertyName: property.key.name, }); } @@ -375,7 +399,21 @@ class Unifier { unify(tA: Type, tB: Type): void { if (tB.kind === "Property") { - const objectType = this.get(tB.object); + if ( + this.env.config.enableTreatRefLikeIdentifiersAsRefs && + isRefLikeName(tB) + ) { + this.unify(tB.objectType, { + kind: "Object", + shapeId: BuiltInUseRefId, + }); + this.unify(tA, { + kind: "Object", + shapeId: BuiltInRefValueId, + }); + return; + } + const objectType = this.get(tB.objectType); const propertyType = this.env.getPropertyType( objectType, tB.propertyName @@ -483,3 +521,9 @@ class Unifier { return type; } } + +const RefLikeNameRE = /^(?:[a-zA-Z$_][a-zA-Z$_0-9]*)Ref$|ref$/; + +function isRefLikeName(t: PropType): boolean { + return RefLikeNameRE.test(t.objectName) && t.propertyName === "current"; +} diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.ref-like-name-not-ref.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.ref-like-name-not-ref.expect.md new file mode 100644 index 0000000000000..b30c22c5b6be4 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.ref-like-name-not-ref.expect.md @@ -0,0 +1,47 @@ + +## Input + +```javascript +// @validatePreserveExistingMemoizationGuarantees +import { useCallback, useRef } from "react"; + +function useCustomRef() { + return useRef({ click: () => {} }); +} + +function Foo() { + const Ref = useCustomRef(); + + const onClick = useCallback(() => { + Ref.current?.click(); + }, []); + + return \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs-2.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs-2.js new file mode 100644 index 0000000000000..3f20e924084cf --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs-2.js @@ -0,0 +1,22 @@ +// @enableTreatRefLikeIdentifiersAsRefs @validatePreserveExistingMemoizationGuarantees +import { useRef, useCallback } from "react"; + +function useCustomRef() { + return useRef({ click: () => {} }); +} + +function Foo() { + const ref = useCustomRef(); + + const onClick = useCallback(() => { + ref.current?.click(); + }, []); + + return \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs.js new file mode 100644 index 0000000000000..07948262b6c1a --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs.js @@ -0,0 +1,22 @@ +// @enableTreatRefLikeIdentifiersAsRefs @validatePreserveExistingMemoizationGuarantees +import { useRef, useCallback } from "react"; + +function useCustomRef() { + return useRef({ click: () => {} }); +} + +function Foo() { + const customRef = useCustomRef(); + + const onClick = useCallback(() => { + customRef.current?.click(); + }, []); + + return