[compiler] Treat ref-like named objects as refs

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<typeof EnvironmentConfigSchema>;

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;

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<IdentifierId, string>,
+  id: IdentifierId,
+  name: Identifier
+): void {
+  if (name.name?.kind === "named") {
+    names.set(id, name.name.value);
+  }
+}
+
+function getName(names: Map<IdentifierId, string>, id: IdentifierId): string {
+  return names.get(id) ?? "";
+}
+
 function* generateInstructionTypes(
   env: Environment,
+  names: Map<IdentifierId, string>,
   instr: Instruction
 ): Generator<TypeEquation, void, undefined> {
   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,
               });
             }
@@ -342,11 +366,11 @@ function* generateInstructionTypes(
       yield equation(left, { kind: "Object", shapeId: BuiltInJsxId });
       break;
     }
+    case "PropertyStore":
     case "DeclareLocal":
     case "NewExpression":
     case "RegExpLiteral":
     case "MetaProperty":
-    case "PropertyStore":
     case "ComputedStore":
     case "ComputedLoad":
     case "TaggedTemplateExpression":
@@ -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";
+}

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 <button onClick={onClick} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+  isComponent: true,
+};
+
+```
+
+
+## Error
+
+```
+   9 |   const Ref = useCustomRef();
+  10 |
+> 11 |   const onClick = useCallback(() => {
+     |                               ^^^^^^^
+> 12 |     Ref.current?.click();
+     | ^^^^^^^^^^^^^^^^^^^^^^^^^
+> 13 |   }, []);
+     | ^^^^ CannotPreserveMemoization: React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. The inferred dependencies did not match the manually specified dependencies, which could cause the value to change more or less frequently than expected (11:13)
+  14 |
+  15 |   return <button onClick={onClick} />;
+  16 | }
+```
+
+

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.ref-like-name-not-Ref.js

@@ -0,0 +1,22 @@
+// @validatePreserveExistingMemoizationGuarantees
+import { useCallback, useRef } from "react";
+
+function useCustomRef() {
+  return useRef({ click: () => {} });
+}
+
+function Foo() {
+  const Ref = useCustomRef();
+
+  const onClick = useCallback(() => {
+    Ref.current?.click();
+  }, []);
+
+  return <button onClick={onClick} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+  isComponent: true,
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.ref-like-name-not-a-ref.expect.md

@@ -0,0 +1,47 @@
+
+## Input
+
+```javascript
+// @validatePreserveExistingMemoizationGuarantees
+import { useCallback, useRef } from "react";
+
+function useCustomRef() {
+  return useRef({ click: () => {} });
+}
+
+function Foo() {
+  const notaref = useCustomRef();
+
+  const onClick = useCallback(() => {
+    notaref.current?.click();
+  }, []);
+
+  return <button onClick={onClick} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+  isComponent: true,
+};
+
+```
+
+
+## Error
+
+```
+   9 |   const notaref = useCustomRef();
+  10 |
+> 11 |   const onClick = useCallback(() => {
+     |                               ^^^^^^^
+> 12 |     notaref.current?.click();
+     | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+> 13 |   }, []);
+     | ^^^^ CannotPreserveMemoization: React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. The inferred dependencies did not match the manually specified dependencies, which could cause the value to change more or less frequently than expected (11:13)
+  14 |
+  15 |   return <button onClick={onClick} />;
+  16 | }
+```
+
+

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.ref-like-name-not-a-ref.js

@@ -0,0 +1,22 @@
+// @validatePreserveExistingMemoizationGuarantees
+import { useCallback, useRef } from "react";
+
+function useCustomRef() {
+  return useRef({ click: () => {} });
+}
+
+function Foo() {
+  const notaref = useCustomRef();
+
+  const onClick = useCallback(() => {
+    notaref.current?.click();
+  }, []);
+
+  return <button onClick={onClick} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+  isComponent: true,
+};