[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: 65bb3289c3756855011cb6f181280287a75eaedf Pull Request resolved: https://github.com/facebook/react/pull/30583
Joe Savona committed
Aug 2, 2024 at 09:53 UTC
1db4d6c41531579fa3a1a8dbcd69d0891d08201f
17 files changed
+272
-84
compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts
+2
-2
@@ -127,11 +127,11 @@ export function* run(
127
code,
128
useMemoCacheIdentifier,
129
);
130
- yield {
130
+ yield log({
131
kind: 'debug',
132
name: 'EnvironmentConfig',
133
value: prettyFormat(env.config),
134
- };
134
+ });
135
const ast = yield* runWithEnvironment(func, env);
136
return ast;
137
}
compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts
+1
@@ -335,6 +335,7 @@ function extractManualMemoizationArgs(
335
export function dropManualMemoization(func: HIRFunction): void {
336
const isValidationEnabled =
337
func.env.config.validatePreserveExistingMemoizationGuarantees ||
338
+ func.env.config.validateNoSetStateInRender ||
339
func.env.config.enablePreserveExistingMemoizationGuarantees;
340
const sidemap: IdentifierSidemap = {
341
functions: new Map(),
compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateNoSetStateInRender.ts
+86
-72
@@ -6,7 +6,7 @@
6
*/
7
8
import {CompilerError, ErrorSeverity} from '../CompilerError';
9
-import {HIRFunction, IdentifierId, Place, isSetStateType} from '../HIR';
9
+import {HIRFunction, IdentifierId, isSetStateType} from '../HIR';
10
import {computeUnconditionalBlocks} from '../HIR/ComputeUnconditionalBlocks';
11
import {eachInstructionValueOperand} from '../HIR/visitors';
12
import {Err, Ok, Result} from '../Utils/Result';
@@ -49,63 +49,97 @@ function validateNoSetStateInRenderImpl(
49
unconditionalSetStateFunctions: Set<IdentifierId>,
50
): Result<void, CompilerError> {
51
const unconditionalBlocks = computeUnconditionalBlocks(fn);
52
-
52
+ let activeManualMemoId: number | null = null;
53
const errors = new CompilerError();
54
for (const [, block] of fn.body.blocks) {
55
- if (unconditionalBlocks.has(block.id)) {
56
- for (const instr of block.instructions) {
57
- switch (instr.value.kind) {
58
- case 'LoadLocal': {
59
- if (
60
- unconditionalSetStateFunctions.has(
61
- instr.value.place.identifier.id,
62
- )
63
- ) {
64
- unconditionalSetStateFunctions.add(instr.lvalue.identifier.id);
65
- }
66
- break;
55
+ for (const instr of block.instructions) {
56
+ switch (instr.value.kind) {
57
+ case 'LoadLocal': {
58
+ if (
59
+ unconditionalSetStateFunctions.has(instr.value.place.identifier.id)
60
+ ) {
61
+ unconditionalSetStateFunctions.add(instr.lvalue.identifier.id);
62
}
68
- case 'StoreLocal': {
69
- if (
70
- unconditionalSetStateFunctions.has(
71
- instr.value.value.identifier.id,
72
- )
73
- ) {
74
- unconditionalSetStateFunctions.add(
75
- instr.value.lvalue.place.identifier.id,
76
- );
77
- unconditionalSetStateFunctions.add(instr.lvalue.identifier.id);
78
- }
79
- break;
80
- }
81
- case 'ObjectMethod':
82
- case 'FunctionExpression': {
83
- if (
84
- // faster-path to check if the function expression references a setState
85
- [...eachInstructionValueOperand(instr.value)].some(
86
- operand =>
87
- isSetStateType(operand.identifier) ||
88
- unconditionalSetStateFunctions.has(operand.identifier.id),
89
- ) &&
90
- // if yes, does it unconditonally call it?
91
- validateNoSetStateInRenderImpl(
92
- instr.value.loweredFunc.func,
93
- unconditionalSetStateFunctions,
94
- ).isErr()
95
- ) {
96
- // This function expression unconditionally calls a setState
97
- unconditionalSetStateFunctions.add(instr.lvalue.identifier.id);
98
- }
99
- break;
63
+ break;
64
+ }
65
+ case 'StoreLocal': {
66
+ if (
67
+ unconditionalSetStateFunctions.has(instr.value.value.identifier.id)
68
+ ) {
69
+ unconditionalSetStateFunctions.add(
70
+ instr.value.lvalue.place.identifier.id,
71
+ );
72
+ unconditionalSetStateFunctions.add(instr.lvalue.identifier.id);
73
}
101
- case 'CallExpression': {
102
- validateNonSetState(
103
- errors,
74
+ break;
75
+ }
76
+ case 'ObjectMethod':
77
+ case 'FunctionExpression': {
78
+ if (
79
+ // faster-path to check if the function expression references a setState
80
+ [...eachInstructionValueOperand(instr.value)].some(
81
+ operand =>
82
+ isSetStateType(operand.identifier) ||
83
+ unconditionalSetStateFunctions.has(operand.identifier.id),
84
+ ) &&
85
+ // if yes, does it unconditonally call it?
86
+ validateNoSetStateInRenderImpl(
87
+ instr.value.loweredFunc.func,
88
unconditionalSetStateFunctions,
105
- instr.value.callee,
106
- );
107
- break;
89
+ ).isErr()
90
+ ) {
91
+ // This function expression unconditionally calls a setState
92
+ unconditionalSetStateFunctions.add(instr.lvalue.identifier.id);
93
}
94
+ break;
95
+ }
96
+ case 'StartMemoize': {
97
+ CompilerError.invariant(activeManualMemoId === null, {
98
+ reason: 'Unexpected nested StartMemoize instructions',
99
+ loc: instr.value.loc,
100
+ });
101
+ activeManualMemoId = instr.value.manualMemoId;
102
+ break;
103
+ }
104
+ case 'FinishMemoize': {
105
+ CompilerError.invariant(
106
+ activeManualMemoId === instr.value.manualMemoId,
107
+ {
108
+ reason:
109
+ 'Expected FinishMemoize to align with previous StartMemoize instruction',
110
+ loc: instr.value.loc,
111
+ },
112
+ );
113
+ activeManualMemoId = null;
114
+ break;
115
+ }
116
+ case 'CallExpression': {
117
+ const callee = instr.value.callee;
118
+ if (
119
+ isSetStateType(callee.identifier) ||
120
+ unconditionalSetStateFunctions.has(callee.identifier.id)
121
+ ) {
122
+ if (activeManualMemoId !== null) {
123
+ errors.push({
124
+ reason:
125
+ 'Calling setState from useMemo may trigger an infinite loop. (https://react.dev/reference/react/useState)',
126
+ description: null,
127
+ severity: ErrorSeverity.InvalidReact,
128
+ loc: callee.loc,
129
+ suggestions: null,
130
+ });
131
+ } else if (unconditionalBlocks.has(block.id)) {
132
+ errors.push({
133
+ reason:
134
+ 'This is an unconditional set state during render, which will trigger an infinite loop. (https://react.dev/reference/react/useState)',
135
+ description: null,
136
+ severity: ErrorSeverity.InvalidReact,
137
+ loc: callee.loc,
138
+ suggestions: null,
139
+ });
140
+ }
141
+ }
142
+ break;
143
}
144
}
145
}
@@ -117,23 +151,3 @@ function validateNoSetStateInRenderImpl(
151
return Ok(undefined);
152
}
153
}
120
-
121
-function validateNonSetState(
122
- errors: CompilerError,
123
- unconditionalSetStateFunctions: Set<IdentifierId>,
124
- operand: Place,
125
-): void {
126
- if (
127
- isSetStateType(operand.identifier) ||
128
- unconditionalSetStateFunctions.has(operand.identifier.id)
129
- ) {
130
- errors.push({
131
- reason:
132
- 'This is an unconditional set state during render, which will trigger an infinite loop. (https://react.dev/reference/react/useState)',
133
- description: null,
134
- severity: ErrorSeverity.InvalidReact,
135
- loc: typeof operand.loc !== 'symbol' ? operand.loc : null,
136
- suggestions: null,
137
- });
138
- }
139
-}
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-global-mutation-unused-usecallback.expect.md
+3
@@ -36,6 +36,9 @@ function Component() {
36
}
37
return t0;
38
}
39
+function _temp() {
40
+ window.foo = true;
41
+}
42
43
export const FIXTURE_ENTRYPOINT = {
44
fn: Component,
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-access-in-effect-indirect.expect.md
+2
-2
@@ -2,7 +2,7 @@
2
## Input
3
4
```javascript
5
-// @validateRefAccessDuringRender
5
+// @validateRefAccessDuringRender @validateNoSetStateInRender:false
6
import {useCallback, useEffect, useRef, useState} from 'react';
7
8
function Component() {
@@ -42,7 +42,7 @@ export const FIXTURE_ENTRYPOINT = {
42
## Code
43
44
```javascript
45
-import { c as _c } from "react/compiler-runtime"; // @validateRefAccessDuringRender
45
+import { c as _c } from "react/compiler-runtime"; // @validateRefAccessDuringRender @validateNoSetStateInRender:false
46
import { useCallback, useEffect, useRef, useState } from "react";
47
48
function Component() {
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-access-in-effect-indirect.js
+1
-1
@@ -1,4 +1,4 @@
1
-// @validateRefAccessDuringRender
1
+// @validateRefAccessDuringRender @validateNoSetStateInRender:false
2
import {useCallback, useEffect, useRef, useState} from 'react';
3
4
function Component() {
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-conditional-setState-in-useMemo.expect.md
new
+36
@@ -0,0 +1,36 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+function Component({item, cond}) {
6
+ const [prevItem, setPrevItem] = useState(item);
7
+ const [state, setState] = useState(0);
8
+
9
+ useMemo(() => {
10
+ if (cond) {
11
+ setPrevItem(item);
12
+ setState(0);
13
+ }
14
+ }, [cond, key, init]);
15
+
16
+ return state;
17
+}
18
+
19
+```
20
+
21
+
22
+## Error
23
+
24
+```
25
+ 5 | useMemo(() => {
26
+ 6 | if (cond) {
27
+> 7 | setPrevItem(item);
28
+ | ^^^^^^^^^^^ InvalidReact: Calling setState from useMemo may trigger an infinite loop. (https://react.dev/reference/react/useState) (7:7)
29
+
30
+InvalidReact: Calling setState from useMemo may trigger an infinite loop. (https://react.dev/reference/react/useState) (8:8)
31
+ 8 | setState(0);
32
+ 9 | }
33
+ 10 | }, [cond, key, init]);
34
+```
35
+
36
+
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-conditional-setState-in-useMemo.js
new
+13
@@ -0,0 +1,13 @@
1
+function Component({item, cond}) {
2
+ const [prevItem, setPrevItem] = useState(item);
3
+ const [state, setState] = useState(0);
4
+
5
+ useMemo(() => {
6
+ if (cond) {
7
+ setPrevItem(item);
8
+ setState(0);
9
+ }
10
+ }, [cond, key, init]);
11
+
12
+ return state;
13
+}
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo-indirect-useCallback.expect.md
new
+38
@@ -0,0 +1,38 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+import {useCallback} from 'react';
6
+
7
+function useKeyedState({key, init}) {
8
+ const [prevKey, setPrevKey] = useState(key);
9
+ const [state, setState] = useState(init);
10
+
11
+ const fn = useCallback(() => {
12
+ setPrevKey(key);
13
+ setState(init);
14
+ });
15
+
16
+ useMemo(() => {
17
+ fn();
18
+ }, [key, init]);
19
+
20
+ return state;
21
+}
22
+
23
+```
24
+
25
+
26
+## Error
27
+
28
+```
29
+ 11 |
30
+ 12 | useMemo(() => {
31
+> 13 | fn();
32
+ | ^^ InvalidReact: Calling setState from useMemo may trigger an infinite loop. (https://react.dev/reference/react/useState) (13:13)
33
+ 14 | }, [key, init]);
34
+ 15 |
35
+ 16 | return state;
36
+```
37
+
38
+
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo-indirect-useCallback.js
new
+17
@@ -0,0 +1,17 @@
1
+import {useCallback} from 'react';
2
+
3
+function useKeyedState({key, init}) {
4
+ const [prevKey, setPrevKey] = useState(key);
5
+ const [state, setState] = useState(init);
6
+
7
+ const fn = useCallback(() => {
8
+ setPrevKey(key);
9
+ setState(init);
10
+ });
11
+
12
+ useMemo(() => {
13
+ fn();
14
+ }, [key, init]);
15
+
16
+ return state;
17
+}
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo.expect.md
new
+34
@@ -0,0 +1,34 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+function useKeyedState({key, init}) {
6
+ const [prevKey, setPrevKey] = useState(key);
7
+ const [state, setState] = useState(init);
8
+
9
+ useMemo(() => {
10
+ setPrevKey(key);
11
+ setState(init);
12
+ }, [key, init]);
13
+
14
+ return state;
15
+}
16
+
17
+```
18
+
19
+
20
+## Error
21
+
22
+```
23
+ 4 |
24
+ 5 | useMemo(() => {
25
+> 6 | setPrevKey(key);
26
+ | ^^^^^^^^^^ InvalidReact: Calling setState from useMemo may trigger an infinite loop. (https://react.dev/reference/react/useState) (6:6)
27
+
28
+InvalidReact: Calling setState from useMemo may trigger an infinite loop. (https://react.dev/reference/react/useState) (7:7)
29
+ 7 | setState(init);
30
+ 8 | }, [key, init]);
31
+ 9 |
32
+```
33
+
34
+
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-useMemo.js
new
+11
@@ -0,0 +1,11 @@
1
+function useKeyedState({key, init}) {
2
+ const [prevKey, setPrevKey] = useState(key);
3
+ const [state, setState] = useState(init);
4
+
5
+ useMemo(() => {
6
+ setPrevKey(key);
7
+ setState(init);
8
+ }, [key, init]);
9
+
10
+ return state;
11
+}
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro.expect.md
+20
-6
@@ -24,12 +24,14 @@ function Component(props) {
24
```javascript
25
import { c as _c } from "react/compiler-runtime";
26
function Component(props) {
27
- const $ = _c(2);
27
+ const $ = _c(7);
28
const item = props.item;
29
let t0;
30
+ let baseVideos;
31
+ let thumbnails;
32
if ($[0] !== item) {
31
- const thumbnails = [];
32
- const baseVideos = getBaseVideos(item);
33
+ thumbnails = [];
34
+ baseVideos = getBaseVideos(item);
35
36
baseVideos.forEach((video) => {
37
const baseVideo = video.hasBaseVideo;
@@ -37,14 +39,26 @@ function Component(props) {
39
thumbnails.push({ extraVideo: true });
40
}
41
});
40
-
41
- t0 = <FlatList baseVideos={baseVideos} items={thumbnails} />;
42
$[0] = item;
43
$[1] = t0;
44
+ $[2] = baseVideos;
45
+ $[3] = thumbnails;
46
} else {
47
t0 = $[1];
48
+ baseVideos = $[2];
49
+ thumbnails = $[3];
50
+ }
51
+ t0 = undefined;
52
+ let t1;
53
+ if ($[4] !== baseVideos || $[5] !== thumbnails) {
54
+ t1 = <FlatList baseVideos={baseVideos} items={thumbnails} />;
55
+ $[4] = baseVideos;
56
+ $[5] = thumbnails;
57
+ $[6] = t1;
58
+ } else {
59
+ t1 = $[6];
60
}
47
- return t0;
61
+ return t1;
62
}
63
64
```
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-named-function.expect.md
+2
-1
@@ -2,6 +2,7 @@
2
## Input
3
4
```javascript
5
+// @validateNoSetStateInRender:false
6
import {useMemo} from 'react';
7
import {makeArray} from 'shared-runtime';
8
@@ -20,7 +21,7 @@ export const FIXTURE_ENTRYPOINT = {
21
## Code
22
23
```javascript
23
-import { c as _c } from "react/compiler-runtime";
24
+import { c as _c } from "react/compiler-runtime"; // @validateNoSetStateInRender:false
25
import { useMemo } from "react";
26
import { makeArray } from "shared-runtime";
27
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-named-function.ts
+1
@@ -1,3 +1,4 @@
1
+// @validateNoSetStateInRender:false
2
import {useMemo} from 'react';
3
import {makeArray} from 'shared-runtime';
4
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-nested-ifs.expect.md
+2
@@ -24,10 +24,12 @@ export const FIXTURE_ENTRYPOINT = {
24
25
```javascript
26
function Component(props) {
27
+ let t0;
28
if (props.cond) {
29
if (props.cond) {
30
}
31
}
32
+ t0 = undefined;
33
}
34
35
export const FIXTURE_ENTRYPOINT = {
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-return-empty.expect.md
+3
@@ -15,7 +15,10 @@ function component(a) {
15
16
```javascript
17
function component(a) {
18
+ let t0;
19
+
20
mutate(a);
21
+ t0 = undefined;
22
}
23
24
```