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

Visitor extension for transforming ReactiveFunction

Adds a subclass of ReactiveFunctionVisitor, ReactiveFunctionTransform, which makes it easier to write passes that change the shape of a ReactiveFunction. The two use-cases converted so far are both flattening away certain categories of reactive scopes — this will make it easier to add a similar pass to prune scopes that contain hook calls.

Joe Savona committed Feb 8, 2023 at 14:07 UTC c7e3bc4d413fd5a17597bd345dc8766ceed983ea
3 files changed +156 -157
compiler/forget/src/ReactiveScopes/FlattenReactiveLoops.ts
+40 -116
@@ -6,142 +6,66 @@
6 */
7
8 import {
9 - ReactiveBlock,
9 ReactiveFunction,
10 ReactiveScopeBlock,
11 + ReactiveStatement,
12 + ReactiveTerminal,
13 + ReactiveTerminalStatement,
14 } from "../HIR/HIR";
15 import { assertExhaustive } from "../Utils/utils";
16 +import {
17 + ReactiveFunctionTransform,
18 + Transformed,
19 + visitReactiveFunction,
20 +} from "./visitors";
21
22 /**
23 * Given a reactive function, flattens any scopes contained within a loop construct.
24 * We won't initially support memoization within loops though this is possible in the future.
25 */
26 export function flattenReactiveLoops(fn: ReactiveFunction): void {
20 - visit(fn.body, false);
27 + visitReactiveFunction(fn, new Transform(), false);
28 }
29
23 -function visit(block: ReactiveBlock, shouldFlatten: boolean): void {
24 - let i = 0;
25 - while (i < block.length) {
26 - const item = block[i]!;
27 - switch (item.kind) {
28 - case "scope": {
29 - if (shouldFlatten) {
30 - const successors = block.splice(i + 1);
31 - block.pop(); // remove the current element
32 - flatten(item, block);
33 - i = block.length;
34 - block.push(...successors);
35 - } else {
36 - visit(item.instructions, false);
37 - i++;
38 - }
39 - break;
40 - }
41 - case "instruction": {
42 - i++;
43 - break;
44 - }
45 - case "terminal": {
46 - const terminal = item.terminal;
47 - switch (terminal.kind) {
48 - case "break":
49 - case "continue":
50 - case "return":
51 - case "throw": {
52 - break;
53 - }
54 - case "for": {
55 - visit(terminal.loop, true);
56 - break;
57 - }
58 - case "while": {
59 - visit(terminal.loop, true);
60 - break;
61 - }
62 - case "if": {
63 - visit(terminal.consequent, shouldFlatten);
64 - if (terminal.alternate !== null) {
65 - visit(terminal.alternate, shouldFlatten);
66 - }
67 - break;
68 - }
69 - case "switch": {
70 - for (const case_ of terminal.cases) {
71 - if (case_.block !== undefined) {
72 - visit(case_.block, shouldFlatten);
73 - }
74 - }
75 - break;
76 - }
77 - default: {
78 - assertExhaustive(
79 - terminal,
80 - `Unexpected terminal kind '${(terminal as any).kind}'`
81 - );
82 - }
83 - }
84 - i++;
85 - break;
86 - }
87 - default: {
88 - assertExhaustive(item, `Unexpected item`);
89 - }
30 +class Transform extends ReactiveFunctionTransform<boolean> {
31 + override transformScope(
32 + scope: ReactiveScopeBlock,
33 + isWithinLoop: boolean
34 + ): Transformed<ReactiveStatement> {
35 + this.visitScope(scope, isWithinLoop);
36 + if (isWithinLoop) {
37 + return { kind: "replace-many", value: scope.instructions };
38 + } else {
39 + return { kind: "keep" };
40 }
41 }
92 -}
42
94 -function flatten(scope: ReactiveScopeBlock, block: ReactiveBlock): void {
95 - for (const item of scope.instructions) {
96 - switch (item.kind) {
97 - case "scope": {
98 - flatten(item, block);
43 + override visitTerminal(
44 + stmt: ReactiveTerminalStatement<ReactiveTerminal>,
45 + isWithinLoop: boolean
46 + ): void {
47 + switch (stmt.terminal.kind) {
48 + // Loop terminals flatten nested scopes
49 + case "while":
50 + case "for": {
51 + this.traverseTerminal(stmt, true);
52 break;
53 }
101 - case "terminal": {
102 - const terminal = item.terminal;
103 - switch (terminal.kind) {
104 - case "break":
105 - case "continue":
106 - case "return":
107 - case "throw": {
108 - break;
109 - }
110 - case "for": {
111 - visit(terminal.loop, true);
112 - break;
113 - }
114 - case "while": {
115 - visit(terminal.loop, true);
116 - break;
117 - }
118 - case "if": {
119 - visit(terminal.consequent, true);
120 - if (terminal.alternate !== null) {
121 - visit(terminal.alternate, true);
122 - }
123 - break;
124 - }
125 - case "switch": {
126 - for (const case_ of terminal.cases) {
127 - if (case_.block !== undefined) {
128 - visit(case_.block, true);
129 - }
130 - }
131 - break;
132 - }
133 - default: {
134 - assertExhaustive(
135 - terminal,
136 - `Unexpected terminal kind '${(terminal as any).kind}'`
137 - );
138 - }
139 - }
140 - block.push(item);
54 + // Non-loop terminals passthrough is contextual, inherits the parent isWithinScope
55 + case "break":
56 + case "continue":
57 + case "if":
58 + case "return":
59 + case "switch":
60 + case "throw": {
61 + this.traverseTerminal(stmt, isWithinLoop);
62 break;
63 }
64 default: {
144 - block.push(item);
65 + assertExhaustive(
66 + stmt.terminal,
67 + `Unexpected terminal kind '${(stmt.terminal as any).kind}'`
68 + );
69 }
70 }
71 }
compiler/forget/src/ReactiveScopes/PruneUnusedScopes.ts
+25 -41
@@ -5,54 +5,38 @@
5 * LICENSE file in the root directory of this source tree.
6 */
7
8 -import { ReactiveBlock, ReactiveFunction } from "../HIR/HIR";
9 -import { assertExhaustive } from "../Utils/utils";
10 -import { mapTerminalBlocks } from "./visitors";
8 +import {
9 + ReactiveFunction,
10 + ReactiveScopeBlock,
11 + ReactiveStatement,
12 +} from "../HIR/HIR";
13 +import {
14 + ReactiveFunctionTransform,
15 + Transformed,
16 + visitReactiveFunction,
17 +} from "./visitors";
18
19 /**
20 * Converts scopes without outputs into regular blocks.
21 */
22 export function pruneUnusedScopes(fn: ReactiveFunction): void {
16 - fn.body = visitBlock(fn.body);
23 + visitReactiveFunction(fn, new Transform(), undefined);
24 }
25
19 -function visitBlock(block: ReactiveBlock): ReactiveBlock {
20 - let nextBlock: ReactiveBlock | null = null;
21 - for (let i = 0; i < block.length; i++) {
22 - const stmt = block[i]!;
23 - switch (stmt.kind) {
24 - case "terminal": {
25 - mapTerminalBlocks(stmt.terminal, visitBlock);
26 - break;
27 - }
28 - case "instruction": {
29 - break;
30 - }
31 - case "scope": {
32 - stmt.instructions = visitBlock(stmt.instructions);
33 - // If a scope doesn't have declarations but reassigns a value, the scope shouldn't be pruned
34 - // as we still want to generate a memo block for that scope
35 - if (
36 - stmt.scope.declarations.size === 0 &&
37 - (stmt.scope.dependencies.size === 0 ||
38 - stmt.scope.reassignments.size === 0)
39 - ) {
40 - nextBlock ??= block.slice(0, i);
41 - nextBlock.push(...stmt.instructions);
42 - continue;
43 - }
44 - break;
45 - }
46 - default: {
47 - assertExhaustive(
48 - stmt,
49 - `Unexpected statement kind '${(stmt as any).kind}'`
50 - );
51 - }
52 - }
53 - if (nextBlock !== null) {
54 - nextBlock.push(stmt);
26 +class Transform extends ReactiveFunctionTransform<void> {
27 + override transformScope(
28 + scopeBlock: ReactiveScopeBlock,
29 + state: void
30 + ): Transformed<ReactiveStatement> {
31 + this.visitScope(scopeBlock, state);
32 + if (
33 + scopeBlock.scope.declarations.size === 0 &&
34 + (scopeBlock.scope.dependencies.size === 0 ||
35 + scopeBlock.scope.reassignments.size === 0)
36 + ) {
37 + return { kind: "replace-many", value: scopeBlock.instructions };
38 + } else {
39 + return { kind: "keep" };
40 }
41 }
57 - return nextBlock ?? block;
42 }
compiler/forget/src/ReactiveScopes/visitors.ts
+91
@@ -13,6 +13,7 @@ import {
13 ReactiveFunction,
14 ReactiveInstruction,
15 ReactiveScopeBlock,
16 + ReactiveStatement,
17 ReactiveTerminal,
18 ReactiveTerminalStatement,
19 ReactiveValue,
@@ -175,6 +176,96 @@ export class ReactiveFunctionVisitor<TState = void> {
176 }
177 }
178
179 +export type Transformed<T> =
180 + | { kind: "remove" }
181 + | { kind: "keep" }
182 + | { kind: "replace"; value: T }
183 + | { kind: "replace-many"; value: Array<T> };
184 +
185 +export class ReactiveFunctionTransform<
186 + TState = void
187 +> extends ReactiveFunctionVisitor<TState> {
188 + override traverseBlock(block: ReactiveBlock, state: TState): void {
189 + let nextBlock: ReactiveBlock | null = null;
190 + for (let i = 0; i < block.length; i++) {
191 + const instr = block[i]!;
192 + let transformed: Transformed<ReactiveStatement>;
193 + switch (instr.kind) {
194 + case "instruction": {
195 + transformed = this.transformInstruction(instr.instruction, state);
196 + break;
197 + }
198 + case "scope": {
199 + transformed = this.transformScope(instr, state);
200 + break;
201 + }
202 + case "terminal": {
203 + transformed = this.transformTerminal(instr, state);
204 + break;
205 + }
206 + default: {
207 + assertExhaustive(
208 + instr,
209 + `Unexpected instruction kind '${(instr as any).kind}'`
210 + );
211 + }
212 + }
213 + switch (transformed.kind) {
214 + case "keep": {
215 + if (nextBlock !== null) {
216 + nextBlock.push(instr);
217 + }
218 + break;
219 + }
220 + case "remove": {
221 + if (nextBlock === null) {
222 + nextBlock = block.slice(0, i);
223 + }
224 + break;
225 + }
226 + case "replace": {
227 + nextBlock ??= block.slice(0, i);
228 + nextBlock.push(transformed.value);
229 + break;
230 + }
231 + case "replace-many": {
232 + nextBlock ??= block.slice(0, i);
233 + nextBlock.push(...transformed.value);
234 + break;
235 + }
236 + }
237 + }
238 + if (nextBlock !== null) {
239 + block.length = 0;
240 + block.push(...nextBlock);
241 + }
242 + }
243 +
244 + transformInstruction(
245 + instruction: ReactiveInstruction,
246 + state: TState
247 + ): Transformed<ReactiveStatement> {
248 + this.visitInstruction(instruction, state);
249 + return { kind: "keep" };
250 + }
251 +
252 + transformTerminal(
253 + stmt: ReactiveTerminalStatement,
254 + state: TState
255 + ): Transformed<ReactiveStatement> {
256 + this.visitTerminal(stmt, state);
257 + return { kind: "keep" };
258 + }
259 +
260 + transformScope(
261 + scope: ReactiveScopeBlock,
262 + state: TState
263 + ): Transformed<ReactiveStatement> {
264 + this.visitScope(scope, state);
265 + return { kind: "keep" };
266 + }
267 +}
268 +
269 export function* eachReactiveValueOperand(
270 instrValue: ReactiveValue
271 ): Iterable<Place> {