From 00fb5c661402650ff993a5cef95b7a904cd9c237 Mon Sep 17 00:00:00 2001
From: Sathya Gunsasekaran <gsathya.ceg@gmail.com>
Date: Mon, 17 Jun 2024 16:07:03 +0100
Subject: [PATCH] [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.
---
 .../src/HIR/Environment.ts                    | 17 ++++
 .../src/HIR/Types.ts                          | 12 ++-
 .../src/TypeInference/InferTypes.ts           | 54 ++++++++++--
 .../error.ref-like-name-not-ref.expect.md     | 47 +++++++++++
 .../compiler/error.ref-like-name-not-ref.js   | 22 +++++
 .../ref-like-name-as-refs-2.expect.md         | 81 ++++++++++++++++++
 .../compiler/ref-like-name-as-refs-2.js       | 22 +++++
 .../ref-like-name-as-refs-3.expect.md         | 84 +++++++++++++++++++
 .../compiler/ref-like-name-as-refs-3.js       | 22 +++++
 .../compiler/ref-like-name-as-refs.expect.md  | 81 ++++++++++++++++++
 .../compiler/ref-like-name-as-refs.js         | 22 +++++
 11 files changed, 455 insertions(+), 9 deletions(-)
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.ref-like-name-not-ref.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.ref-like-name-not-ref.js
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs-2.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs-2.js
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs-3.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs-3.js
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs.js

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<typeof EnvironmentConfigSchema>;
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<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,
               });
             }
@@ -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 <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 | }
+```
+          
+      
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.ref-like-name-not-ref.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.ref-like-name-not-ref.js
new file mode 100644
index 0000000000000..c2386fcbcf3f5
--- /dev/null
+++ b/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,
+};
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs-2.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs-2.expect.md
new file mode 100644
index 0000000000000..855cfc2cd1ded
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs-2.expect.md
@@ -0,0 +1,81 @@
+
+## Input
+
+```javascript
+// @enableTreatRefLikeIdentifiersAsRefs @validatePreserveExistingMemoizationGuarantees
+import { useRef, useCallback } 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,
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enableTreatRefLikeIdentifiersAsRefs @validatePreserveExistingMemoizationGuarantees
+import { useRef, useCallback } from "react";
+
+function useCustomRef() {
+  const $ = _c(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = { click: () => {} };
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return useRef(t0);
+}
+
+function Foo() {
+  const $ = _c(3);
+  const ref = useCustomRef();
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = () => {
+      ref.current?.click();
+    };
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const onClick = t0;
+  let t1;
+  if ($[1] !== onClick) {
+    t1 = <button onClick={onClick} />;
+    $[1] = onClick;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+  isComponent: true,
+};
+
+```
+      
+### Eval output
+(kind: ok) <button></button>
\ 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 <button onClick={onClick} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+  isComponent: true,
+};
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs-3.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs-3.expect.md
new file mode 100644
index 0000000000000..7ebf84dd5ef6c
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs-3.expect.md
@@ -0,0 +1,84 @@
+
+## Input
+
+```javascript
+// @enableTreatRefLikeIdentifiersAsRefs @validatePreserveExistingMemoizationGuarantees
+import { useRef, useEffect } from "react";
+
+function useCustomRef() {
+  return useRef({ click: () => {} });
+}
+
+function Foo() {
+  const ref = useCustomRef();
+
+  useEffect(() => {
+    ref.current?.click();
+  }, []);
+
+  return <div>foo</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+  isComponent: true,
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enableTreatRefLikeIdentifiersAsRefs @validatePreserveExistingMemoizationGuarantees
+import { useRef, useEffect } from "react";
+
+function useCustomRef() {
+  const $ = _c(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = { click: () => {} };
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return useRef(t0);
+}
+
+function Foo() {
+  const $ = _c(3);
+  const ref = useCustomRef();
+  let t0;
+  let t1;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = () => {
+      ref.current?.click();
+    };
+    t1 = [];
+    $[0] = t0;
+    $[1] = t1;
+  } else {
+    t0 = $[0];
+    t1 = $[1];
+  }
+  useEffect(t0, t1);
+  let t2;
+  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+    t2 = <div>foo</div>;
+    $[2] = t2;
+  } else {
+    t2 = $[2];
+  }
+  return t2;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+  isComponent: true,
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>foo</div>
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs-3.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs-3.js
new file mode 100644
index 0000000000000..3e854fa60c21a
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs-3.js
@@ -0,0 +1,22 @@
+// @enableTreatRefLikeIdentifiersAsRefs @validatePreserveExistingMemoizationGuarantees
+import { useRef, useEffect } from "react";
+
+function useCustomRef() {
+  return useRef({ click: () => {} });
+}
+
+function Foo() {
+  const ref = useCustomRef();
+
+  useEffect(() => {
+    ref.current?.click();
+  }, []);
+
+  return <div>foo</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+  isComponent: true,
+};
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs.expect.md
new file mode 100644
index 0000000000000..2df670c8df7a8
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/ref-like-name-as-refs.expect.md
@@ -0,0 +1,81 @@
+
+## Input
+
+```javascript
+// @enableTreatRefLikeIdentifiersAsRefs @validatePreserveExistingMemoizationGuarantees
+import { useRef, useCallback } from "react";
+
+function useCustomRef() {
+  return useRef({ click: () => {} });
+}
+
+function Foo() {
+  const customRef = useCustomRef();
+
+  const onClick = useCallback(() => {
+    customRef.current?.click();
+  }, []);
+
+  return <button onClick={onClick} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+  isComponent: true,
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enableTreatRefLikeIdentifiersAsRefs @validatePreserveExistingMemoizationGuarantees
+import { useRef, useCallback } from "react";
+
+function useCustomRef() {
+  const $ = _c(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = { click: () => {} };
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return useRef(t0);
+}
+
+function Foo() {
+  const $ = _c(3);
+  const customRef = useCustomRef();
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = () => {
+      customRef.current?.click();
+    };
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const onClick = t0;
+  let t1;
+  if ($[1] !== onClick) {
+    t1 = <button onClick={onClick} />;
+    $[1] = onClick;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+  isComponent: true,
+};
+
+```
+      
+### Eval output
+(kind: ok) <button></button>
\ 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 <button onClick={onClick} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+  isComponent: true,
+};