[compiler] Validate against setState in useMemo (resubmit of #30552)

ghstack failed to land #30552 properly, resubmitting

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-source-id: 65bb3289c3
Pull Request resolved: https://github.com/facebook/react/pull/30583

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts

@@ -127,11 +127,11 @@ export function* run(
     code,
     useMemoCacheIdentifier,
   );
-  yield {
+  yield log({
     kind: 'debug',
     name: 'EnvironmentConfig',
     value: prettyFormat(env.config),
-  };
+  });
   const ast = yield* runWithEnvironment(func, env);
   return ast;
 }

compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts

@@ -335,6 +335,7 @@ function extractManualMemoizationArgs(
 export function dropManualMemoization(func: HIRFunction): void {
   const isValidationEnabled =
     func.env.config.validatePreserveExistingMemoizationGuarantees ||
+    func.env.config.validateNoSetStateInRender ||
     func.env.config.enablePreserveExistingMemoizationGuarantees;
   const sidemap: IdentifierSidemap = {
     functions: new Map(),

compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateNoSetStateInRender.ts

@@ -6,7 +6,7 @@
  */
 
 import {CompilerError, ErrorSeverity} from '../CompilerError';
-import {HIRFunction, IdentifierId, Place, isSetStateType} from '../HIR';
+import {HIRFunction, IdentifierId, isSetStateType} from '../HIR';
 import {computeUnconditionalBlocks} from '../HIR/ComputeUnconditionalBlocks';
 import {eachInstructionValueOperand} from '../HIR/visitors';
 import {Err, Ok, Result} from '../Utils/Result';
@@ -49,63 +49,97 @@ function validateNoSetStateInRenderImpl(
   unconditionalSetStateFunctions: Set<IdentifierId>,
 ): Result<void, CompilerError> {
   const unconditionalBlocks = computeUnconditionalBlocks(fn);
-
+  let activeManualMemoId: number | null = null;
   const errors = new CompilerError();
   for (const [, block] of fn.body.blocks) {
-    if (unconditionalBlocks.has(block.id)) {
-      for (const instr of block.instructions) {
-        switch (instr.value.kind) {
-          case 'LoadLocal': {
-            if (
-              unconditionalSetStateFunctions.has(
-                instr.value.place.identifier.id,
-              )
-            ) {
-              unconditionalSetStateFunctions.add(instr.lvalue.identifier.id);
-            }
-            break;
+    for (const instr of block.instructions) {
+      switch (instr.value.kind) {
+        case 'LoadLocal': {
+          if (
+            unconditionalSetStateFunctions.has(instr.value.place.identifier.id)
+          ) {
+            unconditionalSetStateFunctions.add(instr.lvalue.identifier.id);
           }
-          case 'StoreLocal': {
-            if (
-              unconditionalSetStateFunctions.has(
-                instr.value.value.identifier.id,
-              )
-            ) {
-              unconditionalSetStateFunctions.add(
-                instr.value.lvalue.place.identifier.id,
-              );
-              unconditionalSetStateFunctions.add(instr.lvalue.identifier.id);
-            }
-            break;
-          }
-          case 'ObjectMethod':
-          case 'FunctionExpression': {
-            if (
-              // faster-path to check if the function expression references a setState
-              [...eachInstructionValueOperand(instr.value)].some(
-                operand =>
-                  isSetStateType(operand.identifier) ||
-                  unconditionalSetStateFunctions.has(operand.identifier.id),
-              ) &&
-              // if yes, does it unconditonally call it?
-              validateNoSetStateInRenderImpl(
-                instr.value.loweredFunc.func,
-                unconditionalSetStateFunctions,
-              ).isErr()
-            ) {
-              // This function expression unconditionally calls a setState
-              unconditionalSetStateFunctions.add(instr.lvalue.identifier.id);
-            }
-            break;
-          }
-          case 'CallExpression': {
-            validateNonSetState(
-              errors,
-              unconditionalSetStateFunctions,
-              instr.value.callee,
+          break;
+        }
+        case 'StoreLocal': {
+          if (
+            unconditionalSetStateFunctions.has(instr.value.value.identifier.id)
+          ) {
+            unconditionalSetStateFunctions.add(
+              instr.value.lvalue.place.identifier.id,
             );
-            break;
+            unconditionalSetStateFunctions.add(instr.lvalue.identifier.id);
           }
+          break;
+        }
+        case 'ObjectMethod':
+        case 'FunctionExpression': {
+          if (
+            // faster-path to check if the function expression references a setState
+            [...eachInstructionValueOperand(instr.value)].some(
+              operand =>
+                isSetStateType(operand.identifier) ||
+                unconditionalSetStateFunctions.has(operand.identifier.id),
+            ) &&
+            // if yes, does it unconditonally call it?
+            validateNoSetStateInRenderImpl(
+              instr.value.loweredFunc.func,
+              unconditionalSetStateFunctions,
+            ).isErr()
+          ) {
+            // This function expression unconditionally calls a setState
+            unconditionalSetStateFunctions.add(instr.lvalue.identifier.id);
+          }
+          break;
+        }
+        case 'StartMemoize': {
+          CompilerError.invariant(activeManualMemoId === null, {
+            reason: 'Unexpected nested StartMemoize instructions',
+            loc: instr.value.loc,
+          });
+          activeManualMemoId = instr.value.manualMemoId;
+          break;
+        }
+        case 'FinishMemoize': {
+          CompilerError.invariant(
+            activeManualMemoId === instr.value.manualMemoId,
+            {
+              reason:
+                'Expected FinishMemoize to align with previous StartMemoize instruction',
+              loc: instr.value.loc,
+            },
+          );
+          activeManualMemoId = null;
+          break;
+        }
+        case 'CallExpression': {
+          const callee = instr.value.callee;
+          if (
+            isSetStateType(callee.identifier) ||
+            unconditionalSetStateFunctions.has(callee.identifier.id)
+          ) {
+            if (activeManualMemoId !== null) {
+              errors.push({
+                reason:
+                  'Calling setState from useMemo may trigger an infinite loop. (https://react.dev/reference/react/useState)',
+                description: null,
+                severity: ErrorSeverity.InvalidReact,
+                loc: callee.loc,
+                suggestions: null,
+              });
+            } else if (unconditionalBlocks.has(block.id)) {
+              errors.push({
+                reason:
+                  'This is an unconditional set state during render, which will trigger an infinite loop. (https://react.dev/reference/react/useState)',
+                description: null,
+                severity: ErrorSeverity.InvalidReact,
+                loc: callee.loc,
+                suggestions: null,
+              });
+            }
+          }
+          break;
         }
       }
     }
@@ -117,23 +151,3 @@ function validateNoSetStateInRenderImpl(
     return Ok(undefined);
   }
 }
-
-function validateNonSetState(
-  errors: CompilerError,
-  unconditionalSetStateFunctions: Set<IdentifierId>,
-  operand: Place,
-): void {
-  if (
-    isSetStateType(operand.identifier) ||
-    unconditionalSetStateFunctions.has(operand.identifier.id)
-  ) {
-    errors.push({
-      reason:
-        'This is an unconditional set state during render, which will trigger an infinite loop. (https://react.dev/reference/react/useState)',
-      description: null,
-      severity: ErrorSeverity.InvalidReact,
-      loc: typeof operand.loc !== 'symbol' ? operand.loc : null,
-      suggestions: null,
-    });
-  }
-}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-global-mutation-unused-usecallback.expect.md

@@ -36,6 +36,9 @@ function Component() {
   }
   return t0;
 }
+function _temp() {
+  window.foo = true;
+}
 
 export const FIXTURE_ENTRYPOINT = {
   fn: Component,

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-access-in-effect-indirect.expect.md

@@ -2,7 +2,7 @@
 ## Input
 
 ```javascript
-// @validateRefAccessDuringRender
+// @validateRefAccessDuringRender @validateNoSetStateInRender:false
 import {useCallback, useEffect, useRef, useState} from 'react';
 
 function Component() {
@@ -42,7 +42,7 @@ export const FIXTURE_ENTRYPOINT = {
 ## Code
 
 ```javascript
-import { c as _c } from "react/compiler-runtime"; // @validateRefAccessDuringRender
+import { c as _c } from "react/compiler-runtime"; // @validateRefAccessDuringRender @validateNoSetStateInRender:false
 import { useCallback, useEffect, useRef, useState } from "react";
 
 function Component() {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-access-in-effect-indirect.js

@@ -1,4 +1,4 @@
-// @validateRefAccessDuringRender
+// @validateRefAccessDuringRender @validateNoSetStateInRender:false
 import {useCallback, useEffect, useRef, useState} from 'react';
 
 function Component() {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-conditional-setState-in-useMemo.expect.md

@@ -0,0 +1,36 @@
+
+## Input
+
+```javascript
+function Component({item, cond}) {
+  const [prevItem, setPrevItem] = useState(item);
+  const [state, setState] = useState(0);
+
+  useMemo(() => {
+    if (cond) {
+      setPrevItem(item);
+      setState(0);
+    }
+  }, [cond, key, init]);
+
+  return state;
+}
+
+```
+
+
+## Error
+
+```
+   5 |   useMemo(() => {
+   6 |     if (cond) {
+>  7 |       setPrevItem(item);
+     |       ^^^^^^^^^^^ InvalidReact: Calling setState from useMemo may trigger an infinite loop. (https://react.dev/reference/react/useState) (7:7)
+
+InvalidReact: Calling setState from useMemo may trigger an infinite loop. (https://react.dev/reference/react/useState) (8:8)
+   8 |       setState(0);
+   9 |     }
+  10 |   }, [cond, key, init]);
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-conditional-setState-in-useMemo.js

@@ -0,0 +1,13 @@
+function Component({item, cond}) {
+  const [prevItem, setPrevItem] = useState(item);
+  const [state, setState] = useState(0);
+
+  useMemo(() => {
+    if (cond) {
+      setPrevItem(item);
+      setState(0);
+    }
+  }, [cond, key, init]);
+
+  return state;
+}

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;
+```
+          
+      

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

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo.expect.md

@@ -0,0 +1,34 @@
+
+## Input
+
+```javascript
+function useKeyedState({key, init}) {
+  const [prevKey, setPrevKey] = useState(key);
+  const [state, setState] = useState(init);
+
+  useMemo(() => {
+    setPrevKey(key);
+    setState(init);
+  }, [key, init]);
+
+  return state;
+}
+
+```
+
+
+## Error
+
+```
+  4 |
+  5 |   useMemo(() => {
+> 6 |     setPrevKey(key);
+    |     ^^^^^^^^^^ InvalidReact: Calling setState from useMemo may trigger an infinite loop. (https://react.dev/reference/react/useState) (6:6)
+
+InvalidReact: Calling setState from useMemo may trigger an infinite loop. (https://react.dev/reference/react/useState) (7:7)
+  7 |     setState(init);
+  8 |   }, [key, init]);
+  9 |
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo.js

@@ -0,0 +1,11 @@
+function useKeyedState({key, init}) {
+  const [prevKey, setPrevKey] = useState(key);
+  const [state, setState] = useState(init);
+
+  useMemo(() => {
+    setPrevKey(key);
+    setState(init);
+  }, [key, init]);
+
+  return state;
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro.expect.md

@@ -24,12 +24,14 @@ function Component(props) {
 ```javascript
 import { c as _c } from "react/compiler-runtime";
 function Component(props) {
-  const $ = _c(2);
+  const $ = _c(7);
   const item = props.item;
   let t0;
+  let baseVideos;
+  let thumbnails;
   if ($[0] !== item) {
-    const thumbnails = [];
-    const baseVideos = getBaseVideos(item);
+    thumbnails = [];
+    baseVideos = getBaseVideos(item);
 
     baseVideos.forEach((video) => {
       const baseVideo = video.hasBaseVideo;
@@ -37,14 +39,26 @@ function Component(props) {
         thumbnails.push({ extraVideo: true });
       }
     });
-
-    t0 = <FlatList baseVideos={baseVideos} items={thumbnails} />;
     $[0] = item;
     $[1] = t0;
+    $[2] = baseVideos;
+    $[3] = thumbnails;
   } else {
     t0 = $[1];
+    baseVideos = $[2];
+    thumbnails = $[3];
   }
-  return t0;
+  t0 = undefined;
+  let t1;
+  if ($[4] !== baseVideos || $[5] !== thumbnails) {
+    t1 = <FlatList baseVideos={baseVideos} items={thumbnails} />;
+    $[4] = baseVideos;
+    $[5] = thumbnails;
+    $[6] = t1;
+  } else {
+    t1 = $[6];
+  }
+  return t1;
 }
 
 ```

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-named-function.expect.md

@@ -2,6 +2,7 @@
 ## Input
 
 ```javascript
+// @validateNoSetStateInRender:false
 import {useMemo} from 'react';
 import {makeArray} from 'shared-runtime';
 
@@ -20,7 +21,7 @@ export const FIXTURE_ENTRYPOINT = {
 ## Code
 
 ```javascript
-import { c as _c } from "react/compiler-runtime";
+import { c as _c } from "react/compiler-runtime"; // @validateNoSetStateInRender:false
 import { useMemo } from "react";
 import { makeArray } from "shared-runtime";
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-named-function.ts

@@ -1,3 +1,4 @@
+// @validateNoSetStateInRender:false
 import {useMemo} from 'react';
 import {makeArray} from 'shared-runtime';
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-nested-ifs.expect.md

@@ -24,10 +24,12 @@ export const FIXTURE_ENTRYPOINT = {
 
 ```javascript
 function Component(props) {
+  let t0;
   if (props.cond) {
     if (props.cond) {
     }
   }
+  t0 = undefined;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-return-empty.expect.md

@@ -15,7 +15,10 @@ function component(a) {
 
 ```javascript
 function component(a) {
+  let t0;
+
   mutate(a);
+  t0 = undefined;
 }
 
 ```