From 3764b94118eb032dc4e6327dfa002c3324b12686 Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Thu, 1 Aug 2024 16:52:52 -0700
Subject: [PATCH] Update on "[compiler] Validate against setState in useMemo"

Developers sometimes use `useMemo()` as a way to conditionally execute code, including conditionally calling setState. However, the compiler may remove existing useMemo calls if they are not necessary, which _should_ always be a safe optimization. If the useMemo has side effects (eg sets state), then this isn't safe.

This PR improves ValidateNoSetStateInRender to disallow any setState in useMemo (even if it's conditional), expanding on the previous check for unconditional setState in render. Note that the approach uses the StartMemoize/FinishMemoize instructions added in DropManualMemo to know whether a particular setState call is within a useMemo or not. This means enabling the validation in DropManualMemo when the setState validation is enabled, but that's fine since that validation is on everywhere by default (_except_ for in fixtures, which we have a todo for)

[ghstack-poisoned]
---
 ...-in-useMemo-indirect-useCallback.expect.md | 38 +++++++++++++++++++
 ...etState-in-useMemo-indirect-useCallback.js | 17 +++++++++
 2 files changed, 55 insertions(+)
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo-indirect-useCallback.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo-indirect-useCallback.js

diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo-indirect-useCallback.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo-indirect-useCallback.expect.md
new file mode 100644
index 0000000000..acf84652b2
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo-indirect-useCallback.expect.md
@@ -0,0 +1,38 @@
+
+## Input
+
+```javascript
+import {useCallback} from 'react';
+
+function useKeyedState({key, init}) {
+  const [prevKey, setPrevKey] = useState(key);
+  const [state, setState] = useState(init);
+
+  const fn = useCallback(() => {
+    setPrevKey(key);
+    setState(init);
+  });
+
+  useMemo(() => {
+    fn();
+  }, [key, init]);
+
+  return state;
+}
+
+```
+
+
+## Error
+
+```
+  11 |
+  12 |   useMemo(() => {
+> 13 |     fn();
+     |     ^^ InvalidReact: Calling setState from useMemo may trigger an infinite loop. (https://react.dev/reference/react/useState) (13:13)
+  14 |   }, [key, init]);
+  15 |
+  16 |   return state;
+```
+          
+      
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo-indirect-useCallback.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo-indirect-useCallback.js
new file mode 100644
index 0000000000..dbdf5f5a4d
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo-indirect-useCallback.js
@@ -0,0 +1,17 @@
+import {useCallback} from 'react';
+
+function useKeyedState({key, init}) {
+  const [prevKey, setPrevKey] = useState(key);
+  const [state, setState] = useState(init);
+
+  const fn = useCallback(() => {
+    setPrevKey(key);
+    setState(init);
+  });
+
+  useMemo(() => {
+    fn();
+  }, [key, init]);
+
+  return state;
+}