From 4cbdfd776641bbbfae91e7603593cc5afc52a7af Mon Sep 17 00:00:00 2001
From: Gabor Gyongyi <gaborg1988@gmail.com>
Date: Mon, 3 Mar 2025 21:46:58 +0100
Subject: [PATCH] Check ref access in useState initializer function

---
 .../Validation/ValidateNoRefAccesInRender.ts  |  4 ++-
 .../error.invalid-ref-in-useState.expect.md   | 27 ++++++++++++++++++
 .../compiler/error.invalid-ref-in-useState.js |  6 ++++
 .../__tests__/ReactCompilerRule-test.ts       | 28 +++++++++++++++++++
 4 files changed, 64 insertions(+), 1 deletion(-)
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-ref-in-useState.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-ref-in-useState.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 779207b22e646..b58f4cba4c906 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateNoRefAccesInRender.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateNoRefAccesInRender.ts
@@ -401,7 +401,9 @@ function validateNoRefAccessInRenderImpl(
               }
             }
             for (const operand of eachInstructionValueOperand(instr.value)) {
-              if (hookKind != null) {
+              if (hookKind === 'useState') {
+                validateNoRefValueAccess(errors, env, operand);
+              } else if (hookKind != null) {
                 validateNoDirectRefValueAccess(errors, operand, env);
               } else {
                 validateNoRefAccess(errors, env, operand, operand.loc);
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-ref-in-useState.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-ref-in-useState.expect.md
new file mode 100644
index 0000000000000..ca65b3e1729c0
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-ref-in-useState.expect.md
@@ -0,0 +1,27 @@
+
+## Input
+
+```javascript
+// @validateRefAccessDuringRender
+function Component(props) {
+  const ref = useRef(1);
+  const [state] = useState(() => ref.current);
+  return <div>{state}</div>;
+}
+
+```
+
+
+## Error
+
+```
+  2 | function Component(props) {
+  3 |   const ref = useRef(1);
+> 4 |   const [state] = useState(() => ref.current);
+    |                            ^^^^^^^^^^^^^^^^^ InvalidReact: Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef) (4:4)
+  5 |   return <div>{state}</div>;
+  6 | }
+  7 |
+```
+          
+      
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-ref-in-useState.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-ref-in-useState.js
new file mode 100644
index 0000000000000..3e3367d23011a
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-ref-in-useState.js
@@ -0,0 +1,6 @@
+// @validateRefAccessDuringRender
+function Component(props) {
+  const ref = useRef(1);
+  const [state] = useState(() => ref.current);
+  return <div>{state}</div>;
+}
diff --git a/compiler/packages/eslint-plugin-react-compiler/__tests__/ReactCompilerRule-test.ts b/compiler/packages/eslint-plugin-react-compiler/__tests__/ReactCompilerRule-test.ts
index 71be6b6622eb5..4370350637eef 100644
--- a/compiler/packages/eslint-plugin-react-compiler/__tests__/ReactCompilerRule-test.ts
+++ b/compiler/packages/eslint-plugin-react-compiler/__tests__/ReactCompilerRule-test.ts
@@ -104,6 +104,18 @@ const tests: CompilerTestCases = {
         }
       `,
     },
+    {
+      name: 'Ref access in useEffect hook',
+      code: normalizeIndent`
+        function Component() {
+          const ref = useRef(1);
+          useEffect(() => {
+            ref.current = 2 * ref.current;
+          }, []);
+          return <div>Hello world</div>;
+        }
+      `,
+    },
   ],
   invalid: [
     {
@@ -122,6 +134,22 @@ const tests: CompilerTestCases = {
         },
       ],
     },
+    {
+      name: '[InvalidInput] Ref access in useState initial value function',
+      code: normalizeIndent`
+        function Component(props) {
+          const ref = useRef(1);
+          const [value] = useState(() => ref.current);
+          return value;
+        }
+      `,
+      errors: [
+        {
+          message:
+            'Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef)',
+        },
+      ],
+    },
     {
       name: 'Reportable levels can be configured',
       options: [{reportableLevels: new Set([ErrorSeverity.Todo])}],