From 6a3d6a4382cdafc1260483a6fc5f76593fc038e4 Mon Sep 17 00:00:00 2001
From: Joseph Savona <6425824+josephsavona@users.noreply.github.com>
Date: Fri, 20 Dec 2024 08:56:48 -0800
Subject: [PATCH] [compiler] Allow type cast expressions with refs (#31871)

We report a false positive for the combination of a ref-accessing
function placed inside an array which is they type-cast. Here we teach
ref validation about type casts. I also tried other variants like
`return ref as const` but those already worked.

Closes #31864
---
 .../Validation/ValidateNoRefAccesInRender.ts  |  8 +++
 .../allow-ref-type-cast-in-render.expect.md   | 60 +++++++++++++++++++
 .../compiler/allow-ref-type-cast-in-render.js | 17 ++++++
 3 files changed, 85 insertions(+)
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-type-cast-in-render.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-type-cast-in-render.js

diff --git a/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateNoRefAccesInRender.ts b/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateNoRefAccesInRender.ts
index b361b2016a1dd..4db8c700f387f 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateNoRefAccesInRender.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateNoRefAccesInRender.ts
@@ -305,6 +305,14 @@ function validateNoRefAccessInRenderImpl(
             );
             break;
           }
+          case 'TypeCastExpression': {
+            env.set(
+              instr.lvalue.identifier.id,
+              env.get(instr.value.value.identifier.id) ??
+                refTypeOfType(instr.lvalue),
+            );
+            break;
+          }
           case 'LoadContext':
           case 'LoadLocal': {
             env.set(
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-type-cast-in-render.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-type-cast-in-render.expect.md
new file mode 100644
index 0000000000000..56e3039f63cec
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-type-cast-in-render.expect.md
@@ -0,0 +1,60 @@
+
+## Input
+
+```javascript
+import {useRef} from 'react';
+
+function useArrayOfRef() {
+  const ref = useRef(null);
+  const callback = value => {
+    ref.current = value;
+  };
+  return [callback] as const;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: () => {
+    useArrayOfRef();
+    return 'ok';
+  },
+  params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { useRef } from "react";
+
+function useArrayOfRef() {
+  const $ = _c(1);
+  const ref = useRef(null);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    const callback = (value) => {
+      ref.current = value;
+    };
+
+    t0 = [callback];
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0 as const;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: () => {
+    useArrayOfRef();
+    return "ok";
+  },
+
+  params: [{}],
+};
+
+```
+      
+### Eval output
+(kind: ok) "ok"
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-type-cast-in-render.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-type-cast-in-render.js
new file mode 100644
index 0000000000000..2d0aafeffda4c
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-type-cast-in-render.js
@@ -0,0 +1,17 @@
+import {useRef} from 'react';
+
+function useArrayOfRef() {
+  const ref = useRef(null);
+  const callback = value => {
+    ref.current = value;
+  };
+  return [callback] as const;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: () => {
+    useArrayOfRef();
+    return 'ok';
+  },
+  params: [{}],
+};