@samitouri / QOS-React-2 / commits / c69211a9df

[compiler] Prune dependencies that are only used by useRef or useState

Summary: jmbrown215 recently had an observation that the arguments to useState/useRef are only used when a component renders for the first time, and never afterwards. We can skip more computation that we previously could, with reactive blocks that previously recomputed values when inputs changed now only ever computing them on the first render. ghstack-source-id: 5d044ef787a7da901c70990f4399aa90c9b96802 Pull Request resolved: https://github.com/facebook/react/pull/29653

Mike Vitousek committed May 31, 2024 at 14:06 UTC c69211a9dfa683038b1a758aba2ca09c7862a6d3
8 files changed +524 -14
compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts
+10
@@ -91,6 +91,7 @@ import {
91 validatePreservedManualMemoization,
92 validateUseMemo,
93 } from "../Validation";
94 +import pruneInitializationDependencies from "../ReactiveScopes/PruneInitializationDependencies";
95
96 export type CompilerPipelineValue =
97 | { kind: "ast"; name: string; value: CodegenFunction }
@@ -379,6 +380,15 @@ function* runWithEnvironment(
380 value: reactiveFunction,
381 });
382
383 + if (env.config.enableChangeDetectionForDebugging != null) {
384 + pruneInitializationDependencies(reactiveFunction);
385 + yield log({
386 + kind: "reactive",
387 + name: "PruneInitializationDependencies",
388 + value: reactiveFunction,
389 + });
390 + }
391 +
392 propagateEarlyReturns(reactiveFunction);
393 yield log({
394 kind: "reactive",
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneInitializationDependencies.ts new
+290
@@ -0,0 +1,290 @@
1 +/**
2 + * Copyright (c) Meta Platforms, Inc. and affiliates.
3 + *
4 + * This source code is licensed under the MIT license found in the
5 + * LICENSE file in the root directory of this source tree.
6 + */
7 +
8 +import { CompilerError } from "../CompilerError";
9 +import {
10 + Environment,
11 + Identifier,
12 + IdentifierId,
13 + InstructionId,
14 + Place,
15 + ReactiveBlock,
16 + ReactiveFunction,
17 + ReactiveInstruction,
18 + ReactiveScopeBlock,
19 + ReactiveTerminalStatement,
20 + getHookKind,
21 + isUseRefType,
22 + isUseStateType,
23 +} from "../HIR";
24 +import { eachCallArgument, eachInstructionLValue } from "../HIR/visitors";
25 +import DisjointSet from "../Utils/DisjointSet";
26 +import { assertExhaustive } from "../Utils/utils";
27 +import { ReactiveFunctionVisitor, visitReactiveFunction } from "./visitors";
28 +
29 +/**
30 + * This pass is built based on the observation by @jbrown215 that arguments
31 + * to useState and useRef are only used the first time a component is rendered.
32 + * Any subsequent times, the arguments will be evaluated but ignored. In this pass,
33 + * we use this fact to improve the output of the compiler by not recomputing values that
34 + * are only used as arguments (or inputs to arguments to) useState and useRef.
35 + *
36 + * This pass isn't yet stress-tested so it's not enabled by default. It's only enabled
37 + * to support certain debug modes that detect non-idempotent code, since non-idempotent
38 + * code can "safely" be used if its only passed to useState and useRef. We plan to rewrite
39 + * this pass in HIR and enable it as an optimization in the future.
40 + *
41 + * Algorithm:
42 + * We take two passes over the reactive function AST. In the first pass, we gather
43 + * aliases and build relationships between property accesses--the key thing we need
44 + * to do here is to find that, e.g., $0.x and $1 refer to the same value if
45 + * $1 = PropertyLoad $0.x.
46 + *
47 + * In the second pass, we traverse the AST in reverse order and track how each place
48 + * is used. If a place is read from in any Terminal, we mark the place as "Update", meaning
49 + * it is used whenever the component is updated/re-rendered. If a place is read from in
50 + * a useState or useRef hook call, we mark it as "Create", since it is only used when the
51 + * component is created. In other instructions, we propagate the inferred place for the
52 + * instructions lvalues onto any other instructions that are read.
53 + *
54 + * Whenever we finish this reverse pass over a reactive block, we can look at the blocks
55 + * dependencies and see whether the dependencies are used in an "Update" context or only
56 + * in a "Create" context. If a dependency is create-only, then we can remove that dependency
57 + * from the block.
58 + */
59 +
60 +type CreateUpdate = "Create" | "Update" | "Unknown";
61 +
62 +type KindMap = Map<IdentifierId, CreateUpdate>;
63 +
64 +class Visitor extends ReactiveFunctionVisitor<CreateUpdate> {
65 + map: KindMap = new Map();
66 + aliases: DisjointSet<IdentifierId>;
67 + paths: Map<IdentifierId, Map<string, IdentifierId>>;
68 + env: Environment;
69 +
70 + constructor(
71 + env: Environment,
72 + aliases: DisjointSet<IdentifierId>,
73 + paths: Map<IdentifierId, Map<string, IdentifierId>>
74 + ) {
75 + super();
76 + this.aliases = aliases;
77 + this.paths = paths;
78 + this.env = env;
79 + }
80 +
81 + join(values: Array<CreateUpdate>): CreateUpdate {
82 + function join2(l: CreateUpdate, r: CreateUpdate): CreateUpdate {
83 + if (l === "Update" || r === "Update") {
84 + return "Update";
85 + } else if (l === "Create" || r === "Create") {
86 + return "Create";
87 + } else if (l === "Unknown" || r === "Unknown") {
88 + return "Unknown";
89 + }
90 + assertExhaustive(r, `Unhandled variable kind ${r}`);
91 + }
92 + return values.reduce(join2, "Unknown");
93 + }
94 +
95 + isCreateOnlyHook(id: Identifier): boolean {
96 + return isUseStateType(id) || isUseRefType(id);
97 + }
98 +
99 + override visitPlace(
100 + _: InstructionId,
101 + place: Place,
102 + state: CreateUpdate
103 + ): void {
104 + this.map.set(
105 + place.identifier.id,
106 + this.join([state, this.map.get(place.identifier.id) ?? "Unknown"])
107 + );
108 + }
109 +
110 + override visitBlock(block: ReactiveBlock, state: CreateUpdate): void {
111 + super.visitBlock([...block].reverse(), state);
112 + }
113 +
114 + override visitInstruction(instruction: ReactiveInstruction): void {
115 + const state = this.join(
116 + [...eachInstructionLValue(instruction)].map(
117 + (operand) => this.map.get(operand.identifier.id) ?? "Unknown"
118 + )
119 + );
120 +
121 + const visitCallOrMethodNonArgs = (): void => {
122 + switch (instruction.value.kind) {
123 + case "CallExpression": {
124 + this.visitPlace(instruction.id, instruction.value.callee, state);
125 + break;
126 + }
127 + case "MethodCall": {
128 + this.visitPlace(instruction.id, instruction.value.property, state);
129 + this.visitPlace(instruction.id, instruction.value.receiver, state);
130 + break;
131 + }
132 + }
133 + };
134 +
135 + const isHook = (): boolean => {
136 + let callee = null;
137 + switch (instruction.value.kind) {
138 + case "CallExpression": {
139 + callee = instruction.value.callee.identifier;
140 + break;
141 + }
142 + case "MethodCall": {
143 + callee = instruction.value.property.identifier;
144 + break;
145 + }
146 + }
147 + return callee != null && getHookKind(this.env, callee) != null;
148 + };
149 +
150 + switch (instruction.value.kind) {
151 + case "CallExpression":
152 + case "MethodCall": {
153 + if (
154 + instruction.lvalue &&
155 + this.isCreateOnlyHook(instruction.lvalue.identifier)
156 + ) {
157 + [...eachCallArgument(instruction.value.args)].forEach((operand) =>
158 + this.visitPlace(instruction.id, operand, "Create")
159 + );
160 + visitCallOrMethodNonArgs();
161 + } else {
162 + this.traverseInstruction(instruction, isHook() ? "Update" : state);
163 + }
164 + break;
165 + }
166 + default: {
167 + this.traverseInstruction(instruction, state);
168 + }
169 + }
170 + }
171 +
172 + override visitScope(scope: ReactiveScopeBlock): void {
173 + const state = this.join(
174 + [
175 + ...scope.scope.declarations.keys(),
176 + ...[...scope.scope.reassignments.values()].map((ident) => ident.id),
177 + ].map((id) => this.map.get(id) ?? "Unknown")
178 + );
179 + super.visitScope(scope, state);
180 + [...scope.scope.dependencies].forEach((ident) => {
181 + let target: undefined | IdentifierId =
182 + this.aliases.find(ident.identifier.id) ?? ident.identifier.id;
183 + ident.path.forEach((key) => {
184 + target &&= this.paths.get(target)?.get(key);
185 + });
186 + if (target && this.map.get(target) === "Create") {
187 + scope.scope.dependencies.delete(ident);
188 + }
189 + });
190 + }
191 +
192 + override visitTerminal(
193 + stmt: ReactiveTerminalStatement,
194 + state: CreateUpdate
195 + ): void {
196 + CompilerError.invariant(state !== "Create", {
197 + reason: "Visiting a terminal statement with state 'Create'",
198 + loc: stmt.terminal.loc,
199 + });
200 + super.visitTerminal(stmt, state);
201 + }
202 +
203 + override visitReactiveFunctionValue(
204 + _id: InstructionId,
205 + _dependencies: Array<Place>,
206 + fn: ReactiveFunction,
207 + state: CreateUpdate
208 + ): void {
209 + visitReactiveFunction(fn, this, state);
210 + }
211 +}
212 +
213 +export default function pruneInitializationDependencies(
214 + fn: ReactiveFunction
215 +): void {
216 + const [aliases, paths] = getAliases(fn);
217 + visitReactiveFunction(fn, new Visitor(fn.env, aliases, paths), "Update");
218 +}
219 +
220 +function update(
221 + map: Map<IdentifierId, Map<string, IdentifierId>>,
222 + key: IdentifierId,
223 + path: string,
224 + value: IdentifierId
225 +): void {
226 + const inner = map.get(key) ?? new Map();
227 + inner.set(path, value);
228 + map.set(key, inner);
229 +}
230 +
231 +class AliasVisitor extends ReactiveFunctionVisitor {
232 + scopeIdentifiers: DisjointSet<IdentifierId> = new DisjointSet<IdentifierId>();
233 + scopePaths: Map<IdentifierId, Map<string, IdentifierId>> = new Map();
234 +
235 + override visitInstruction(instr: ReactiveInstruction): void {
236 + if (
237 + instr.value.kind === "StoreLocal" ||
238 + instr.value.kind === "StoreContext"
239 + ) {
240 + this.scopeIdentifiers.union([
241 + instr.value.lvalue.place.identifier.id,
242 + instr.value.value.identifier.id,
243 + ]);
244 + } else if (
245 + instr.value.kind === "LoadLocal" ||
246 + instr.value.kind === "LoadContext"
247 + ) {
248 + instr.lvalue &&
249 + this.scopeIdentifiers.union([
250 + instr.lvalue.identifier.id,
251 + instr.value.place.identifier.id,
252 + ]);
253 + } else if (instr.value.kind === "PropertyLoad") {
254 + instr.lvalue &&
255 + update(
256 + this.scopePaths,
257 + instr.value.object.identifier.id,
258 + instr.value.property,
259 + instr.lvalue.identifier.id
260 + );
261 + } else if (instr.value.kind === "PropertyStore") {
262 + update(
263 + this.scopePaths,
264 + instr.value.object.identifier.id,
265 + instr.value.property,
266 + instr.value.value.identifier.id
267 + );
268 + }
269 + }
270 +}
271 +
272 +function getAliases(
273 + fn: ReactiveFunction
274 +): [DisjointSet<IdentifierId>, Map<IdentifierId, Map<string, IdentifierId>>] {
275 + const visitor = new AliasVisitor();
276 + visitReactiveFunction(fn, visitor, null);
277 + let disjoint = visitor.scopeIdentifiers;
278 + let scopePaths = new Map<IdentifierId, Map<string, IdentifierId>>();
279 + for (const [key, value] of visitor.scopePaths) {
280 + for (const [path, id] of value) {
281 + update(
282 + scopePaths,
283 + disjoint.find(key) ?? key,
284 + path,
285 + disjoint.find(id) ?? id
286 + );
287 + }
288 + }
289 + return [disjoint, scopePaths];
290 +}
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-and-other-hook-unpruned-dependency.expect.md new
+84
@@ -0,0 +1,84 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +import { useState } from "react"; // @enableChangeDetectionForDebugging
6 +
7 +function useOther(x) {
8 + return x;
9 +}
10 +
11 +function Component(props) {
12 + const w = f(props.x);
13 + const z = useOther(w);
14 + const [x, _] = useState(z);
15 + return <div>{x}</div>;
16 +}
17 +
18 +function f(x) {
19 + return x;
20 +}
21 +
22 +export const FIXTURE_ENTRYPOINT = {
23 + fn: Component,
24 + params: [{ x: 42 }],
25 + isComponent: true,
26 +};
27 +
28 +```
29 +
30 +## Code
31 +
32 +```javascript
33 +import { $structuralCheck } from "react-compiler-runtime";
34 +import { c as _c } from "react/compiler-runtime";
35 +import { useState } from "react"; // @enableChangeDetectionForDebugging
36 +
37 +function useOther(x) {
38 + return x;
39 +}
40 +
41 +function Component(props) {
42 + const $ = _c(4);
43 + let t0;
44 + {
45 + t0 = f(props.x);
46 + if (!($[0] !== props.x)) {
47 + let old$t0;
48 + old$t0 = $[1];
49 + $structuralCheck(old$t0, t0, "t0", "Component");
50 + t0 = old$t0;
51 + }
52 + $[0] = props.x;
53 + $[1] = t0;
54 + }
55 + const w = t0;
56 + const z = useOther(w);
57 + const [x] = useState(z);
58 + let t1;
59 + {
60 + t1 = <div>{x}</div>;
61 + if (!($[2] !== x)) {
62 + let old$t1;
63 + old$t1 = $[3];
64 + $structuralCheck(old$t1, t1, "t1", "Component");
65 + t1 = old$t1;
66 + }
67 + $[2] = x;
68 + $[3] = t1;
69 + }
70 + return t1;
71 +}
72 +
73 +function f(x) {
74 + return x;
75 +}
76 +
77 +export const FIXTURE_ENTRYPOINT = {
78 + fn: Component,
79 + params: [{ x: 42 }],
80 + isComponent: true,
81 +};
82 +
83 +```
84 +
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-and-other-hook-unpruned-dependency.js new
+22
@@ -0,0 +1,22 @@
1 +import { useState } from "react"; // @enableChangeDetectionForDebugging
2 +
3 +function useOther(x) {
4 + return x;
5 +}
6 +
7 +function Component(props) {
8 + const w = f(props.x);
9 + const z = useOther(w);
10 + const [x, _] = useState(z);
11 + return <div>{x}</div>;
12 +}
13 +
14 +function f(x) {
15 + return x;
16 +}
17 +
18 +export const FIXTURE_ENTRYPOINT = {
19 + fn: Component,
20 + params: [{ x: 42 }],
21 + isComponent: true,
22 +};
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-pruned-dependency-change-detect.expect.md
+9 -14
@@ -20,31 +20,26 @@ import { c as _c } from "react/compiler-runtime"; // @enableChangeDetectionForDe
20 import { useState } from "react";
21
22 function Component(props) {
23 - const $ = _c(4);
23 + const $ = _c(3);
24 let t0;
25 - {
25 + if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
26 t0 = f(props.x);
27 - if (!($[0] !== props.x)) {
28 - let old$t0;
29 - old$t0 = $[1];
30 - $structuralCheck(old$t0, t0, "t0", "Component");
31 - t0 = old$t0;
32 - }
33 - $[0] = props.x;
34 - $[1] = t0;
27 + $[0] = t0;
28 + } else {
29 + t0 = $[0];
30 }
31 const [x] = useState(t0);
32 let t1;
33 {
34 t1 = <div>{x}</div>;
40 - if (!($[2] !== x)) {
35 + if (!($[1] !== x)) {
36 let old$t1;
42 - old$t1 = $[3];
37 + old$t1 = $[2];
38 $structuralCheck(old$t1, t1, "t1", "Component");
39 t1 = old$t1;
40 }
46 - $[2] = x;
47 - $[3] = t1;
41 + $[1] = x;
42 + $[2] = t1;
43 }
44 return t1;
45 }
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-unpruned-dependency.expect.md new
+85
@@ -0,0 +1,85 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +import { useState } from "react"; // @enableChangeDetectionForDebugging
6 +
7 +function Component(props) {
8 + const w = f(props.x);
9 + const [x, _] = useState(w);
10 + return (
11 + <div>
12 + {x}
13 + {w}
14 + </div>
15 + );
16 +}
17 +
18 +function f(x) {
19 + return x;
20 +}
21 +
22 +export const FIXTURE_ENTRYPOINT = {
23 + fn: Component,
24 + params: [{ x: 42 }],
25 + isComponent: true,
26 +};
27 +
28 +```
29 +
30 +## Code
31 +
32 +```javascript
33 +import { $structuralCheck } from "react-compiler-runtime";
34 +import { c as _c } from "react/compiler-runtime";
35 +import { useState } from "react"; // @enableChangeDetectionForDebugging
36 +
37 +function Component(props) {
38 + const $ = _c(5);
39 + let t0;
40 + {
41 + t0 = f(props.x);
42 + if (!($[0] !== props.x)) {
43 + let old$t0;
44 + old$t0 = $[1];
45 + $structuralCheck(old$t0, t0, "t0", "Component");
46 + t0 = old$t0;
47 + }
48 + $[0] = props.x;
49 + $[1] = t0;
50 + }
51 + const w = t0;
52 + const [x] = useState(w);
53 + let t1;
54 + {
55 + t1 = (
56 + <div>
57 + {x}
58 + {w}
59 + </div>
60 + );
61 + if (!($[2] !== x || $[3] !== w)) {
62 + let old$t1;
63 + old$t1 = $[4];
64 + $structuralCheck(old$t1, t1, "t1", "Component");
65 + t1 = old$t1;
66 + }
67 + $[2] = x;
68 + $[3] = w;
69 + $[4] = t1;
70 + }
71 + return t1;
72 +}
73 +
74 +function f(x) {
75 + return x;
76 +}
77 +
78 +export const FIXTURE_ENTRYPOINT = {
79 + fn: Component,
80 + params: [{ x: 42 }],
81 + isComponent: true,
82 +};
83 +
84 +```
85 +
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-unpruned-dependency.js new
+22
@@ -0,0 +1,22 @@
1 +import { useState } from "react"; // @enableChangeDetectionForDebugging
2 +
3 +function Component(props) {
4 + const w = f(props.x);
5 + const [x, _] = useState(w);
6 + return (
7 + <div>
8 + {x}
9 + {w}
10 + </div>
11 + );
12 +}
13 +
14 +function f(x) {
15 + return x;
16 +}
17 +
18 +export const FIXTURE_ENTRYPOINT = {
19 + fn: Component,
20 + params: [{ x: 42 }],
21 + isComponent: true,
22 +};
compiler/packages/snap/src/SproutTodoFilter.ts
+2
@@ -496,6 +496,8 @@ const skipFilter = new Set([
496
497 "fast-refresh-refresh-on-const-changes-dev",
498 "useState-pruned-dependency-change-detect",
499 + "useState-unpruned-dependency",
500 + "useState-and-other-hook-unpruned-dependency",
501 ]);
502
503 export default skipFilter;