[compiler] Treat ref-like named objects as refs

If a component uses the `useRef` hook directly then we type it's return
value as a ref. But if it's wrapped in a custom hook then we lose out
on this type information as the compiler doesn't look at the hook
definition. This has resulted in some false positives in our analysis
like the ones reported in #29160 and #29196.

This PR 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 = ...;
})
```
In the above example, `ref` and `myRef` will be treated as React refs.

Author: gsathya

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,
+};