@samitouri / QOS-React / commits / 2b2d305199

[hir] Rewrite buildReactiveBlocks -> buildReactiveScopeTerminalsHIR

ghstack-source-id: 1e804dd31d4b5f74070c94c0ea56f55539e46703 Pull Request resolved: https://github.com/facebook/react-forget/pull/2853

Mofei Zhang committed Apr 29, 2024 at 14:08 UTC 2b2d30519991f0d956a37d149a3947cca6784ec1
13 files changed +384 -43
compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts
+18 -9
@@ -16,6 +16,7 @@ import {
16 assertTerminalSuccessorsExist,
17 assertValidBlockNesting,
18 assertValidMutableRanges,
19 + buildReactiveScopeTerminalsHIR,
20 lower,
21 mergeConsecutiveBlocks,
22 mergeOverlappingReactiveScopesHIR,
@@ -238,7 +239,7 @@ function* runWithEnvironment(
239 value: hir,
240 });
241
241 - if (env.config.enableAlignReactiveScopesToBlockScopesHIR) {
242 + if (env.config.enableReactiveScopesInHIR) {
243 pruneUnusedLabelsHIR(hir);
244 yield log({
245 kind: "hir",
@@ -259,6 +260,14 @@ function* runWithEnvironment(
260 name: "MergeOverlappingReactiveScopesHIR",
261 value: hir,
262 });
263 + assertValidBlockNesting(hir);
264 +
265 + buildReactiveScopeTerminalsHIR(hir);
266 + yield log({
267 + kind: "hir",
268 + name: "BuildReactiveScopeTerminalsHIR",
269 + value: hir,
270 + });
271
272 assertValidBlockNesting(hir);
273 }
@@ -277,7 +286,7 @@ function* runWithEnvironment(
286 value: reactiveFunction,
287 });
288
280 - if (!env.config.enableAlignReactiveScopesToBlockScopesHIR) {
289 + if (!env.config.enableReactiveScopesInHIR) {
290 alignReactiveScopesToBlockScopes(reactiveFunction);
291 yield log({
292 kind: "reactive",
@@ -291,14 +300,14 @@ function* runWithEnvironment(
300 name: "MergeOverlappingReactiveScopes",
301 value: reactiveFunction,
302 });
294 - }
303
296 - buildReactiveBlocks(reactiveFunction);
297 - yield log({
298 - kind: "reactive",
299 - name: "BuildReactiveBlocks",
300 - value: reactiveFunction,
301 - });
304 + buildReactiveBlocks(reactiveFunction);
305 + yield log({
306 + kind: "reactive",
307 + name: "BuildReactiveBlocks",
308 + value: reactiveFunction,
309 + });
310 + }
311
312 flattenReactiveLoops(reactiveFunction);
313 yield log({
compiler/packages/babel-plugin-react-forget/src/HIR/AssertValidBlockNesting.ts
+48 -24
@@ -49,7 +49,7 @@ type Block =
49 id: ScopeId;
50 } & MutableRange);
51
52 -function getScopes(fn: HIRFunction): Set<ReactiveScope> {
52 +export function getScopes(fn: HIRFunction): Set<ReactiveScope> {
53 const scopes: Set<ReactiveScope> = new Set();
54 function visitPlace(place: Place): void {
55 const scope = place.identifier.scope;
@@ -94,12 +94,57 @@ function getScopes(fn: HIRFunction): Set<ReactiveScope> {
94 * 6 |
95 * 7 ⌟
96 */
97 -function nestedRangeComparator(a: MutableRange, b: MutableRange): number {
97 +export function rangePreOrderComparator(
98 + a: MutableRange,
99 + b: MutableRange
100 +): number {
101 const startDiff = a.start - b.start;
102 if (startDiff !== 0) return startDiff;
103 return b.end - a.end;
104 }
105
106 +export function recursivelyTraverseItems<T, TContext>(
107 + items: Array<T>,
108 + getRange: (val: T) => MutableRange,
109 + context: TContext,
110 + enter: (val: T, context: TContext) => void,
111 + exit: (val: T, context: TContext) => void
112 +): void {
113 + items.sort((a, b) => rangePreOrderComparator(getRange(a), getRange(b)));
114 + let activeItems: Array<T> = [];
115 + const ranges = items.map(getRange);
116 + for (let i = 0; i < items.length; i++) {
117 + const curr = items[i];
118 + const currRange = ranges[i];
119 + for (let i = activeItems.length - 1; i >= 0; i--) {
120 + const maybeParent = activeItems[i];
121 + const maybeParentRange = getRange(maybeParent);
122 + const disjoint = currRange.start >= maybeParentRange.end;
123 + const nested = currRange.end <= maybeParentRange.end;
124 + CompilerError.invariant(disjoint || nested, {
125 + reason: "Invalid nesting in program blocks or scopes",
126 + description: `Items overlap but are not nested: ${maybeParentRange.start}:${maybeParentRange.end}(${currRange.start}:${currRange.end})`,
127 + loc: GeneratedSource,
128 + });
129 + if (disjoint) {
130 + exit(maybeParent, context);
131 + activeItems.length = i;
132 + } else {
133 + break;
134 + }
135 + }
136 + enter(curr, context);
137 + activeItems.push(curr);
138 + }
139 +
140 + let curr = activeItems.pop();
141 + while (curr != null) {
142 + exit(curr, context);
143 + curr = activeItems.pop();
144 + }
145 +}
146 +const no_op: () => void = () => {};
147 +
148 export function assertValidBlockNesting(fn: HIRFunction): void {
149 const scopes = getScopes(fn);
150
@@ -122,26 +167,5 @@ export function assertValidBlockNesting(fn: HIRFunction): void {
167 }
168 }
169
125 - blocks.sort(nestedRangeComparator);
126 -
127 - let active: Array<Block> = [];
128 - for (let i = 0; i < blocks.length; i++) {
129 - const curr = blocks[i];
130 - for (let i = active.length - 1; i >= 0; i--) {
131 - const maybeParent = active[i];
132 - const disjoint = curr.start >= maybeParent.end;
133 - const nested = curr.end <= maybeParent.end;
134 - CompilerError.invariant(disjoint || nested, {
135 - reason: "Invalid nesting in program blocks or scopes",
136 - description: `Blocks overlap but are not nested: ${maybeParent.kind}@${maybeParent.id}(${maybeParent.start}:${maybeParent.end}) ${curr.kind}@${curr.id}(${curr.start}:${curr.end})`,
137 - loc: GeneratedSource,
138 - });
139 - if (disjoint) {
140 - active.length = i;
141 - } else {
142 - break;
143 - }
144 - }
145 - active.push(curr);
146 - }
170 + recursivelyTraverseItems(blocks, (block) => block, null, no_op, no_op);
171 }
compiler/packages/babel-plugin-react-forget/src/HIR/BuildReactiveScopeTerminalsHIR.ts new
+288
@@ -0,0 +1,288 @@
1 +import { CompilerError } from "../CompilerError";
2 +import { getScopes, recursivelyTraverseItems } from "./AssertValidBlockNesting";
3 +import { Environment } from "./Environment";
4 +import {
5 + BasicBlock,
6 + BlockId,
7 + GeneratedSource,
8 + GotoTerminal,
9 + GotoVariant,
10 + HIRFunction,
11 + InstructionId,
12 + ReactiveScope,
13 + ReactiveScopeTerminal,
14 + ScopeId,
15 +} from "./HIR";
16 +
17 +/**
18 + * This pass assumes that all program blocks are properly nested with respect to fallthroughs
19 + * (e.g. a valid javascript AST).
20 + * Given a function whose reactive scope ranges have been correctly aligned and merged,
21 + * this pass rewrites blocks to introduce ReactiveScopeTerminals and their fallthrough blocks.
22 + * e.g.
23 + * ```js
24 + * // source
25 + * [0] ...
26 + * [1] const x = []; ⌝ scope range
27 + * [2] if (cond) { |
28 + * [3] x.push(a); |
29 + * } |
30 + * [4] x.push(b); ⌟
31 + * [5] ...
32 + *
33 + * // before this pass
34 + * bb0:
35 + * [0]
36 + * [1]
37 + * [2]
38 + * If ($2) then bb1 else bb2 (fallthrough=bb2)
39 + * bb1:
40 + * [3]
41 + * Goto bb2
42 + * bb2:
43 + * [4]
44 + * [5]
45 + *
46 + * // after this pass
47 + * bb0:
48 + * [0]
49 + * ScopeTerminal goto=bb3 (fallthrough=bb4) <-- new
50 + * bb3: <-- new
51 + * [1]
52 + * [2]
53 + * If ($2) then bb1 else bb2 (fallthrough=bb2)
54 + * bb1:
55 + * [3]
56 + * Goto bb2
57 + * bb2:
58 + * [4]
59 + * Goto bb4 <-- new
60 + * bb4: <-- new
61 + * [5]
62 + * ```
63 + */
64 +
65 +export function buildReactiveScopeTerminalsHIR(fn: HIRFunction): void {
66 + /**
67 + * Step 1:
68 + * Traverse all blocks to build up a list of rewrites. We also pre-allocate the
69 + * fallthrough ID here as scope start terminals and scope end terminals both
70 + * require a fallthrough block.
71 + */
72 + const queuedRewrites: Array<TerminalRewriteInfo> = [];
73 + recursivelyTraverseItems(
74 + [...getScopes(fn)],
75 + (scope) => scope.range,
76 + {
77 + fallthroughs: new Map(),
78 + rewrites: queuedRewrites,
79 + env: fn.env,
80 + },
81 + pushStartScopeTerminal,
82 + pushEndScopeTerminal
83 + );
84 +
85 + /**
86 + * Step 2:
87 + * Traverse all blocks to apply rewrites. Here, we split blocks as described at
88 + * the top of this file to add scope terminals and fallthroughs.
89 + */
90 + const rewrittenFinalBlocks = new Map<BlockId, BlockId>();
91 + const nextBlocks = new Map<BlockId, BasicBlock>();
92 + /**
93 + * reverse queuedRewrites to pop off the end as we traverse instructions in
94 + * ascending order
95 + */
96 + queuedRewrites.reverse();
97 + for (const [, block] of fn.body.blocks) {
98 + const context: RewriteContext = {
99 + nextBlockId: block.id,
100 + rewrites: [],
101 + nextPreds: block.preds,
102 + instrSliceIdx: 0,
103 + source: block,
104 + };
105 + /**
106 + * Handle queued terminal rewrites at their nearest instruction ID.
107 + * Note that multiple terminal rewrites may map to the same instruction ID.
108 + */
109 + for (let i = 0; i < block.instructions.length + 1; i++) {
110 + const instrId =
111 + i < block.instructions.length
112 + ? block.instructions[i].id
113 + : block.terminal.id;
114 + let rewrite = queuedRewrites.at(-1);
115 + while (rewrite != null && rewrite.instrId <= instrId) {
116 + handleRewrite(rewrite, i, context);
117 + queuedRewrites.pop();
118 + rewrite = queuedRewrites.at(-1);
119 + }
120 + }
121 +
122 + if (context.rewrites.length > 0) {
123 + const finalBlock: BasicBlock = {
124 + id: context.nextBlockId,
125 + kind: block.kind,
126 + preds: context.nextPreds,
127 + terminal: block.terminal,
128 + instructions: block.instructions.slice(context.instrSliceIdx),
129 + phis: new Set(),
130 + };
131 + context.rewrites.push(finalBlock);
132 + for (const b of context.rewrites) {
133 + nextBlocks.set(b.id, b);
134 + }
135 + rewrittenFinalBlocks.set(block.id, finalBlock.id);
136 + } else {
137 + nextBlocks.set(block.id, block);
138 + }
139 + }
140 + const originalBlocks = fn.body.blocks;
141 + fn.body.blocks = nextBlocks;
142 +
143 + /**
144 + * Step 3:
145 + * Repoint preds and phis when they refer to a rewritten block.
146 + */
147 + for (const [, block] of originalBlocks) {
148 + for (const pred of block.preds) {
149 + const newId = rewrittenFinalBlocks.get(pred);
150 + if (newId != null) {
151 + block.preds.delete(pred);
152 + block.preds.add(newId);
153 + }
154 + }
155 + for (const phi of block.phis) {
156 + for (const [originalId, value] of phi.operands) {
157 + const newId = rewrittenFinalBlocks.get(originalId);
158 + if (newId != null) {
159 + phi.operands.delete(originalId);
160 + phi.operands.set(newId, value);
161 + }
162 + }
163 + }
164 + }
165 +}
166 +
167 +type TerminalRewriteInfo =
168 + | {
169 + kind: "StartScope";
170 + blockId: BlockId;
171 + fallthroughId: BlockId;
172 + instrId: InstructionId;
173 + scope: ReactiveScope;
174 + }
175 + | {
176 + kind: "EndScope";
177 + instrId: InstructionId;
178 + fallthroughId: BlockId;
179 + };
180 +
181 +/**
182 + * Helpers for reversing scope ranges to gather terminal rewrite information
183 + */
184 +type ScopeTraversalContext = {
185 + // cache allocated fallthroughs for start/end scope terminal pairs
186 + fallthroughs: Map<ScopeId, BlockId>;
187 + rewrites: Array<TerminalRewriteInfo>;
188 + env: Environment;
189 +};
190 +
191 +function pushStartScopeTerminal(
192 + scope: ReactiveScope,
193 + context: ScopeTraversalContext
194 +): void {
195 + const blockId = context.env.nextBlockId;
196 + const fallthroughId = context.env.nextBlockId;
197 + context.rewrites.push({
198 + kind: "StartScope",
199 + blockId,
200 + fallthroughId,
201 + instrId: scope.range.start,
202 + scope,
203 + });
204 + context.fallthroughs.set(scope.id, fallthroughId);
205 +}
206 +
207 +function pushEndScopeTerminal(
208 + scope: ReactiveScope,
209 + context: ScopeTraversalContext
210 +): void {
211 + const fallthroughId = context.fallthroughs.get(scope.id);
212 + CompilerError.invariant(fallthroughId != null, {
213 + reason: "Expected scope to exist",
214 + loc: GeneratedSource,
215 + });
216 + context.rewrites.push({
217 + kind: "EndScope",
218 + fallthroughId,
219 + instrId: scope.range.end,
220 + });
221 +}
222 +
223 +type RewriteContext = {
224 + source: BasicBlock;
225 + instrSliceIdx: number;
226 + nextPreds: Set<BlockId>;
227 + nextBlockId: BlockId;
228 + rewrites: Array<BasicBlock>;
229 +};
230 +
231 +/**
232 + * Create a block rewrite by slicing a set of instructions from source.
233 + * Since scope start-ends always end with a GOTO to the next instruction
234 + * from the source block, we directly connect rewritten blocks using state
235 + * from `context`.
236 + *
237 + * Source:
238 + * bb0:
239 + * instr1, instr2, instr3, instr4, [[ original terminal ]]
240 + * Rewritten:
241 + * bb0:
242 + * instr1, [[ scope start block=bb1]]
243 + * bb1:
244 + * instr2, instr3, [[ scope end goto=bb2 ]]
245 + * bb2:
246 + * instr4, [[ original terminal ]]
247 + */
248 +function handleRewrite(
249 + terminalInfo: TerminalRewriteInfo,
250 + idx: number,
251 + context: RewriteContext
252 +): void {
253 + // TODO make consistent instruction IDs instead of reusing
254 + const terminal: ReactiveScopeTerminal | GotoTerminal =
255 + terminalInfo.kind === "StartScope"
256 + ? {
257 + kind: "scope",
258 + fallthrough: terminalInfo.fallthroughId,
259 + block: terminalInfo.blockId,
260 + scope: terminalInfo.scope,
261 + id: terminalInfo.instrId,
262 + loc: GeneratedSource,
263 + }
264 + : {
265 + kind: "goto",
266 + variant: GotoVariant.Break,
267 + block: terminalInfo.fallthroughId,
268 + id: terminalInfo.instrId,
269 + loc: GeneratedSource,
270 + };
271 +
272 + const currBlockId = context.nextBlockId;
273 + context.rewrites.push({
274 + kind: context.source.kind,
275 + id: currBlockId,
276 + instructions: context.source.instructions.slice(context.instrSliceIdx, idx),
277 + preds: context.nextPreds,
278 + // Only the first rewrite should reuse source block phis
279 + phis: context.rewrites.length === 0 ? context.source.phis : new Set(),
280 + terminal,
281 + });
282 + context.nextPreds = new Set([currBlockId]);
283 + context.nextBlockId =
284 + terminalInfo.kind === "StartScope"
285 + ? terminalInfo.blockId
286 + : terminalInfo.fallthroughId;
287 + context.instrSliceIdx = idx;
288 +}
compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts
+1 -1
@@ -169,7 +169,7 @@ const EnvironmentConfigSchema = z.object({
169 */
170 enableUseTypeAnnotations: z.boolean().default(false),
171
172 - enableAlignReactiveScopesToBlockScopesHIR: z.boolean().default(true),
172 + enableReactiveScopesInHIR: z.boolean().default(true),
173
174 /*
175 * Enable validation of hooks to partially check that the component honors the rules of hooks.
compiler/packages/babel-plugin-react-forget/src/HIR/index.ts
+1
@@ -10,6 +10,7 @@ export { assertTerminalSuccessorsExist } from "./AssertTerminalSuccessorsExist";
10 export { assertValidBlockNesting } from "./AssertValidBlockNesting";
11 export { assertValidMutableRanges } from "./AssertValidMutableRanges";
12 export { lower } from "./BuildHIR";
13 +export { buildReactiveScopeTerminalsHIR } from "./BuildReactiveScopeTerminalsHIR";
14 export { computeDominatorTree, computePostDominatorTree } from "./Dominator";
15 export {
16 Environment,
compiler/packages/babel-plugin-react-forget/src/Utils/utils.ts
+14
@@ -45,6 +45,20 @@ export function retainWhere<T>(
45 array.length = writeIndex;
46 }
47
48 +export function getOrInsertWith<U, V>(
49 + m: Map<U, V>,
50 + key: U,
51 + makeDefault: () => V
52 +): V {
53 + if (m.has(key)) {
54 + return m.get(key) as V;
55 + } else {
56 + const defaultValue = makeDefault();
57 + m.set(key, defaultValue);
58 + return defaultValue;
59 + }
60 +}
61 +
62 export function getOrInsertDefault<U, V>(
63 m: Map<U, V>,
64 key: U,
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.bug-repro-trycatch-nested-overlapping-range.expect.md
+1 -1
@@ -20,7 +20,7 @@ function Foo() {
20 ## Error
21
22 ```
23 -Invariant: Invalid nesting in program blocks or scopes. Blocks overlap but are not nested: Scope@0(2:24) ProgramBlockSubtree@17(18:26)
23 +Invariant: Invalid nesting in program blocks or scopes. Items overlap but are not nested: 2:24(18:26)
24 ```
25
26
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.repro-bug-ref-mutable-range.expect.md
+1 -1
@@ -21,7 +21,7 @@ function Foo(props, ref) {
21 ## Error
22
23 ```
24 -Invariant: Invalid nesting in program blocks or scopes. Blocks overlap but are not nested: Scope@0(1:21) ProgramBlockSubtree@1(16:23)
24 +Invariant: Invalid nesting in program blocks or scopes. Items overlap but are not nested: 1:21(16:23)
25 ```
26
27
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.todo-iife-return-modified-later-logical.expect.md
+1 -1
@@ -23,7 +23,7 @@ export const FIXTURE_ENTRYPOINT = {
23 ## Error
24
25 ```
26 -Invariant: Invalid nesting in program blocks or scopes. Blocks overlap but are not nested: ProgramBlockSubtree@0(2:15) Scope@0(3:21)
26 +Invariant: Invalid nesting in program blocks or scopes. Items overlap but are not nested: 2:15(3:21)
27 ```
28
29
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.todo-reactive-scope-overlaps-if.expect.md
+1 -1
@@ -35,7 +35,7 @@ export const FIXTURE_ENTRYPOINT = {
35 ## Error
36
37 ```
38 -Invariant: Invalid nesting in program blocks or scopes. Blocks overlap but are not nested: ProgramBlockSubtree@2(6:11) Scope@1(7:15)
38 +Invariant: Invalid nesting in program blocks or scopes. Items overlap but are not nested: 6:11(7:15)
39 ```
40
41
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.todo-reactive-scope-overlaps-label.expect.md
+1 -1
@@ -34,7 +34,7 @@ export const FIXTURE_ENTRYPOINT = {
34 ## Error
35
36 ```
37 -Invariant: Invalid nesting in program blocks or scopes. Blocks overlap but are not nested: ProgramBlockSubtree@0(3:14) Scope@0(4:18)
37 +Invariant: Invalid nesting in program blocks or scopes. Items overlap but are not nested: 3:14(4:18)
38 ```
39
40
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.todo-reactive-scope-overlaps-try.expect.md
+1 -1
@@ -30,7 +30,7 @@ export const FIXTURE_ENTRYPOINT = {
30 ## Error
31
32 ```
33 -Invariant: Invalid nesting in program blocks or scopes. Blocks overlap but are not nested: ProgramBlockSubtree@0(4:19) Scope@0(5:22)
33 +Invariant: Invalid nesting in program blocks or scopes. Items overlap but are not nested: 4:19(5:22)
34 ```
35
36
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/sequence-expression.expect.md
+8 -3
@@ -19,7 +19,7 @@ function foo() {}
19 ```javascript
20 import { unstable_useMemoCache as useMemoCache } from "react";
21 function sequence(props) {
22 - const $ = useMemoCache(1);
22 + const $ = useMemoCache(2);
23 let t0;
24 if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
25 t0 = (Math.max(1, 2), foo());
@@ -28,8 +28,13 @@ function sequence(props) {
28 t0 = $[0];
29 }
30 let x = t0;
31 - while ((foo(), true)) {
32 - x = (foo(), 2);
31 + if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
32 + while ((foo(), true)) {
33 + x = (foo(), 2);
34 + }
35 + $[1] = x;
36 + } else {
37 + x = $[1];
38 }
39 return x;
40 }