@samitouri / QOS-React-1 / commits / f4fbeb8803

HIR-based AlignReactiveScopesToBlockScopes

ghstack-source-id: ae560a93f24fdbdb4ed73cdd889feb084bb172e5 Pull Request resolved: https://github.com/facebook/react-forget/pull/2822

Joe Savona committed Apr 8, 2024 at 13:48 UTC f4fbeb88038523edfb37646bc9d5314c0f7423ed
3 files changed +220 -6
compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts
+18 -6
@@ -64,6 +64,7 @@ import {
64 renameVariables,
65 } from "../ReactiveScopes";
66 import { alignMethodCallScopes } from "../ReactiveScopes/AlignMethodCallScopes";
67 +import { alignReactiveScopesToBlockScopesHIR } from "../ReactiveScopes/AlignReactiveScopesToBlockScopesHIR";
68 import { pruneAlwaysInvalidatingScopes } from "../ReactiveScopes/PruneAlwaysInvalidatingScopes";
69 import { eliminateRedundantPhi, enterSSA, leaveSSA } from "../SSA";
70 import { inferTypes } from "../TypeInference";
@@ -233,6 +234,15 @@ function* runWithEnvironment(
234 value: hir,
235 });
236
237 + if (env.config.enableAlignReactiveScopesToBlockScopesHIR) {
238 + alignReactiveScopesToBlockScopesHIR(hir);
239 + yield log({
240 + kind: "hir",
241 + name: "AlignReactiveScopesToBlockScopesHIR",
242 + value: hir,
243 + });
244 + }
245 +
246 const reactiveFunction = buildReactiveFunction(hir);
247 yield log({
248 kind: "reactive",
@@ -247,12 +257,14 @@ function* runWithEnvironment(
257 value: reactiveFunction,
258 });
259
250 - alignReactiveScopesToBlockScopes(reactiveFunction);
251 - yield log({
252 - kind: "reactive",
253 - name: "AlignReactiveScopesToBlockScopes",
254 - value: reactiveFunction,
255 - });
260 + if (!env.config.enableAlignReactiveScopesToBlockScopesHIR) {
261 + alignReactiveScopesToBlockScopes(reactiveFunction);
262 + yield log({
263 + kind: "reactive",
264 + name: "AlignReactiveScopesToBlockScopes",
265 + value: reactiveFunction,
266 + });
267 + }
268
269 mergeOverlappingReactiveScopes(reactiveFunction);
270 yield log({
compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts
+2
@@ -169,6 +169,8 @@ const EnvironmentConfigSchema = z.object({
169 */
170 enableUseTypeAnnotations: z.boolean().default(false),
171
172 + enableAlignReactiveScopesToBlockScopesHIR: z.boolean().default(true),
173 +
174 /*
175 * Enable validation of hooks to partially check that the component honors the rules of hooks.
176 * When disabled, the component is assumed to follow the rules (though the Babel plugin looks
compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/AlignReactiveScopesToBlockScopesHIR.ts new
+200
@@ -0,0 +1,200 @@
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 "..";
9 +import {
10 + BlockId,
11 + HIRFunction,
12 + InstructionId,
13 + Place,
14 + ReactiveScope,
15 + makeInstructionId,
16 +} from "../HIR/HIR";
17 +import {
18 + eachInstructionLValue,
19 + eachInstructionValueOperand,
20 + eachTerminalOperand,
21 + mapTerminalSuccessors,
22 + terminalFallthrough,
23 +} from "../HIR/visitors";
24 +import { retainWhere } from "../Utils/utils";
25 +
26 +/*
27 + * Note: this is the 2nd of 4 passes that determine how to break a function into discrete
28 + * reactive scopes (independently memoizeable units of code):
29 + * 1. InferReactiveScopeVariables (on HIR) determines operands that mutate together and assigns
30 + * them a unique reactive scope.
31 + * 2. AlignReactiveScopesToBlockScopes (this pass, on ReactiveFunction) aligns reactive scopes
32 + * to block scopes.
33 + * 3. MergeOverlappingReactiveScopes (on ReactiveFunction) ensures that reactive scopes do not
34 + * overlap, merging any such scopes.
35 + * 4. BuildReactiveBlocks (on ReactiveFunction) groups the statements for each scope into
36 + * a ReactiveScopeBlock.
37 + *
38 + * Prior inference passes assign a reactive scope to each operand, but the ranges of these
39 + * scopes are based on specific instructions at arbitrary points in the control-flow graph.
40 + * However, to codegen blocks around the instructions in each scope, the scopes must be
41 + * aligned to block-scope boundaries - we can't memoize half of a loop!
42 + *
43 + * This pass updates reactive scope boundaries to align to control flow boundaries, for
44 + * example:
45 + *
46 + * ```javascript
47 + * function foo(cond, a) {
48 + * ⌵ original scope
49 + * ⌵ expanded scope
50 + * const x = []; ⌝ ⌝
51 + * if (cond) { ⎮ ⎮
52 + * ... ⎮ ⎮
53 + * x.push(a); ⌟ ⎮
54 + * ... ⎮
55 + * } ⌟
56 + * }
57 + * ```
58 + *
59 + * Here the original scope for `x` ended partway through the if consequent, but we can't
60 + * memoize part of that block. This pass would align the scope to the end of the consequent.
61 + *
62 + * The more general rule is that a reactive scope may only end at the same block scope as it
63 + * began: this pass therefore finds, for each scope, the block where that scope started and
64 + * finds the first instruction after the scope's mutable range in that same block scope (which
65 + * will be the updated end for that scope).
66 + */
67 +
68 +export function alignReactiveScopesToBlockScopesHIR(fn: HIRFunction): void {
69 + type BlockContext =
70 + | { kind: "block"; block: BlockId; scopes: Array<ReactiveScope> }
71 + | {
72 + kind: "value";
73 + start: InstructionId;
74 + end: InstructionId;
75 + scopes: Array<ReactiveScope>;
76 + };
77 + const blockContexts = new Map<BlockId, BlockContext>();
78 + const seen = new Set<ReactiveScope>();
79 +
80 + function recordPlace(place: Place, context: BlockContext): void {
81 + const scope = place.identifier.scope;
82 + if (scope == null) {
83 + return;
84 + }
85 +
86 + if (seen.has(scope)) {
87 + return;
88 + }
89 + if (context.kind === "value") {
90 + scope.range.start = makeInstructionId(
91 + Math.min(context.start, scope.range.start)
92 + );
93 + scope.range.end = makeInstructionId(
94 + Math.max(context.end, scope.range.end)
95 + );
96 + }
97 + seen.add(scope);
98 + context.scopes.push(scope);
99 + }
100 +
101 + for (const [, block] of fn.body.blocks) {
102 + const { instructions, terminal } = block;
103 + let context = blockContexts.get(block.id);
104 + if (context === undefined) {
105 + if (block.kind === "block" || block.kind === "catch") {
106 + context = { kind: "block", block: block.id, scopes: [] };
107 + } else {
108 + CompilerError.invariant(false, {
109 + reason: `Expected a context to be initialized for value block`,
110 + loc: instructions[0]?.loc ?? terminal.loc,
111 + description: `No value for block bb${block.id}`,
112 + });
113 + }
114 + } else if (block.kind === "block" && context.kind !== "block") {
115 + CompilerError.invariant(false, {
116 + reason: `Expected a block context for block`,
117 + loc: instructions[0]?.loc ?? terminal.loc,
118 + description: `Got value block for bb${block.id}`,
119 + });
120 + }
121 +
122 + /*
123 + * Any scopes that carried over across a terminal->fallback need their range extended
124 + * to at least the first instruction of the fallback
125 + */
126 + const startId = instructions.at(0)?.id ?? terminal.id;
127 + for (const scope of context.scopes) {
128 + scope.range.end = makeInstructionId(Math.max(scope.range.end, startId));
129 + }
130 +
131 + /*
132 + * Visit instructions, pruning scopes that end and recording new scopes that appear
133 + * on operands
134 + */
135 + for (const instr of instructions) {
136 + retainWhere(context.scopes, (scope) => scope.range.end > instr.id);
137 + for (const lvalue of eachInstructionLValue(instr)) {
138 + recordPlace(lvalue, context);
139 + }
140 + for (const operand of eachInstructionValueOperand(instr.value)) {
141 + recordPlace(operand, context);
142 + }
143 + }
144 +
145 + // Close scopes that complete at the terminal, and visit scopes of operands
146 + retainWhere(context.scopes, (scope) => scope.range.end > terminal.id);
147 + for (const operand of eachTerminalOperand(terminal)) {
148 + recordPlace(operand, context);
149 + }
150 +
151 + // Save the current context for the fallback block, where this block scope continues
152 + const fallthrough = terminalFallthrough(terminal);
153 + if (fallthrough !== null && !blockContexts.has(fallthrough)) {
154 + blockContexts.set(fallthrough, context);
155 + }
156 +
157 + /*
158 + * Visit all successors (not just direct successors for control-flow ordering) to
159 + * set a value block context where necessary to align the value block start/end
160 + * back to the outer block scope.
161 + *
162 + * TODO: add a variant of eachTerminalSuccessor() that visits _all_ successors, not
163 + * just those that are direct successors for normal control-flow ordering.
164 + */
165 + mapTerminalSuccessors(terminal, (successor) => {
166 + const successorBlock = fn.body.blocks.get(successor)!;
167 + /*
168 + * we need the block kind check here because the do..while terminal's successor
169 + * is a block, and try's successor is a catch block
170 + */
171 + if (
172 + !blockContexts.has(successor) &&
173 + successorBlock.kind !== "block" &&
174 + successorBlock.kind !== "catch"
175 + ) {
176 + let valueContext: BlockContext;
177 + if (context!.kind === "value") {
178 + valueContext = context!;
179 + } else {
180 + CompilerError.invariant(fallthrough !== null, {
181 + reason: `Expected a fallthrough for value block`,
182 + loc: terminal.loc,
183 + });
184 + const fallthroughBlock = fn.body.blocks.get(fallthrough)!;
185 + const nextId =
186 + fallthroughBlock.instructions[0]?.id ??
187 + fallthroughBlock.terminal.id;
188 + valueContext = {
189 + kind: "value",
190 + start: terminal.id,
191 + end: nextId,
192 + scopes: [],
193 + } as BlockContext;
194 + }
195 + blockContexts.set(successor, valueContext);
196 + }
197 + return successor;
198 + });
199 + }
200 +}