| 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 {getScopes, recursivelyTraverseItems} from './AssertValidBlockNesting'; |
| 10 | import {Environment} from './Environment'; |
| 11 | import { |
| 12 | BasicBlock, |
| 13 | BlockId, |
| 14 | GeneratedSource, |
| 15 | GotoTerminal, |
| 16 | GotoVariant, |
| 17 | HIRFunction, |
| 18 | InstructionId, |
| 19 | ReactiveScope, |
| 20 | ReactiveScopeTerminal, |
| 21 | ScopeId, |
| 22 | } from './HIR'; |
| 23 | import { |
| 24 | fixScopeAndIdentifierRanges, |
| 25 | markInstructionIds, |
| 26 | markPredecessors, |
| 27 | reversePostorderBlocks, |
| 28 | } from './HIRBuilder'; |
| 29 | |
| 30 | /** |
| 31 | * This pass assumes that all program blocks are properly nested with respect to fallthroughs |
| 32 | * (e.g. a valid javascript AST). |
| 33 | * Given a function whose reactive scope ranges have been correctly aligned and merged, |
| 34 | * this pass rewrites blocks to introduce ReactiveScopeTerminals and their fallthrough blocks. |
| 35 | * e.g. |
| 36 | * ```js |
| 37 | * // source |
| 38 | * [0] ... |
| 39 | * [1] const x = []; ⌝ scope range |
| 40 | * [2] if (cond) { | |
| 41 | * [3] x.push(a); | |
| 42 | * } | |
| 43 | * [4] x.push(b); ⌟ |
| 44 | * [5] ... |
| 45 | * |
| 46 | * // before this pass |
| 47 | * bb0: |
| 48 | * [0] |
| 49 | * [1] |
| 50 | * [2] |
| 51 | * If ($2) then bb1 else bb2 (fallthrough=bb2) |
| 52 | * bb1: |
| 53 | * [3] |
| 54 | * Goto bb2 |
| 55 | * bb2: |
| 56 | * [4] |
| 57 | * [5] |
| 58 | * |
| 59 | * // after this pass |
| 60 | * bb0: |
| 61 | * [0] |
| 62 | * ScopeTerminal goto=bb3 (fallthrough=bb4) <-- new |
| 63 | * bb3: <-- new |
| 64 | * [1] |
| 65 | * [2] |
| 66 | * If ($2) then bb1 else bb2 (fallthrough=bb2) |
| 67 | * bb1: |
| 68 | * [3] |
| 69 | * Goto bb2 |
| 70 | * bb2: |
| 71 | * [4] |
| 72 | * Goto bb4 <-- new |
| 73 | * bb4: <-- new |
| 74 | * [5] |
| 75 | * ``` |
| 76 | */ |
| 77 | |
| 78 | export function buildReactiveScopeTerminalsHIR(fn: HIRFunction): void { |
| 79 | /** |
| 80 | * Step 1: |
| 81 | * Traverse all blocks to build up a list of rewrites. We also pre-allocate the |
| 82 | * fallthrough ID here as scope start terminals and scope end terminals both |
| 83 | * require a fallthrough block. |
| 84 | */ |
| 85 | const queuedRewrites: Array<TerminalRewriteInfo> = []; |
| 86 | recursivelyTraverseItems( |
| 87 | [...getScopes(fn)], |
| 88 | scope => scope.range, |
| 89 | { |
| 90 | fallthroughs: new Map(), |
| 91 | rewrites: queuedRewrites, |
| 92 | env: fn.env, |
| 93 | }, |
| 94 | pushStartScopeTerminal, |
| 95 | pushEndScopeTerminal, |
| 96 | ); |
| 97 | |
| 98 | /** |
| 99 | * Step 2: |
| 100 | * Traverse all blocks to apply rewrites. Here, we split blocks as described at |
| 101 | * the top of this file to add scope terminals and fallthroughs. |
| 102 | */ |
| 103 | const rewrittenFinalBlocks = new Map<BlockId, BlockId>(); |
| 104 | const nextBlocks = new Map<BlockId, BasicBlock>(); |
| 105 | /** |
| 106 | * reverse queuedRewrites to pop off the end as we traverse instructions in |
| 107 | * ascending order |
| 108 | */ |
| 109 | queuedRewrites.reverse(); |
| 110 | for (const [, block] of fn.body.blocks) { |
| 111 | const context: RewriteContext = { |
| 112 | nextBlockId: block.id, |
| 113 | rewrites: [], |
| 114 | nextPreds: block.preds, |
| 115 | instrSliceIdx: 0, |
| 116 | source: block, |
| 117 | }; |
| 118 | /** |
| 119 | * Handle queued terminal rewrites at their nearest instruction ID. |
| 120 | * Note that multiple terminal rewrites may map to the same instruction ID. |
| 121 | */ |
| 122 | for (let i = 0; i < block.instructions.length + 1; i++) { |
| 123 | const instrId = |
| 124 | i < block.instructions.length |
| 125 | ? block.instructions[i].id |
| 126 | : block.terminal.id; |
| 127 | let rewrite = queuedRewrites.at(-1); |
| 128 | while (rewrite != null && rewrite.instrId <= instrId) { |
| 129 | handleRewrite(rewrite, i, context); |
| 130 | queuedRewrites.pop(); |
| 131 | rewrite = queuedRewrites.at(-1); |
| 132 | } |
| 133 | } |
| 134 | |
| 135 | if (context.rewrites.length > 0) { |
| 136 | const finalBlock: BasicBlock = { |
| 137 | id: context.nextBlockId, |
| 138 | kind: block.kind, |
| 139 | preds: context.nextPreds, |
| 140 | terminal: block.terminal, |
| 141 | instructions: block.instructions.slice(context.instrSliceIdx), |
| 142 | phis: new Set(), |
| 143 | }; |
| 144 | context.rewrites.push(finalBlock); |
| 145 | for (const b of context.rewrites) { |
| 146 | nextBlocks.set(b.id, b); |
| 147 | } |
| 148 | rewrittenFinalBlocks.set(block.id, finalBlock.id); |
| 149 | } else { |
| 150 | nextBlocks.set(block.id, block); |
| 151 | } |
| 152 | } |
| 153 | const originalBlocks = fn.body.blocks; |
| 154 | fn.body.blocks = nextBlocks; |
| 155 | |
| 156 | /** |
| 157 | * Step 3: |
| 158 | * Repoint phis when they refer to a rewritten block. |
| 159 | */ |
| 160 | for (const [, block] of originalBlocks) { |
| 161 | for (const phi of block.phis) { |
| 162 | for (const [originalId, value] of phi.operands) { |
| 163 | const newId = rewrittenFinalBlocks.get(originalId); |
| 164 | if (newId != null) { |
| 165 | phi.operands.delete(originalId); |
| 166 | phi.operands.set(newId, value); |
| 167 | } |
| 168 | } |
| 169 | } |
| 170 | } |
| 171 | |
| 172 | /** |
| 173 | * Step 4: |
| 174 | * Fixup the HIR to restore RPO, ensure correct predecessors, and |
| 175 | * renumber instructions. Note that the renumbering instructions |
| 176 | * invalidates scope and identifier ranges, so we fix them in the |
| 177 | * next step. |
| 178 | */ |
| 179 | reversePostorderBlocks(fn.body); |
| 180 | markPredecessors(fn.body); |
| 181 | markInstructionIds(fn.body); |
| 182 | |
| 183 | /** |
| 184 | * Step 5: |
| 185 | * Fix scope and identifier ranges to account for renumbered instructions |
| 186 | */ |
| 187 | fixScopeAndIdentifierRanges(fn.body); |
| 188 | } |
| 189 | |
| 190 | type TerminalRewriteInfo = |
| 191 | | { |
| 192 | kind: 'StartScope'; |
| 193 | blockId: BlockId; |
| 194 | fallthroughId: BlockId; |
| 195 | instrId: InstructionId; |
| 196 | scope: ReactiveScope; |
| 197 | } |
| 198 | | { |
| 199 | kind: 'EndScope'; |
| 200 | instrId: InstructionId; |
| 201 | fallthroughId: BlockId; |
| 202 | }; |
| 203 | |
| 204 | /** |
| 205 | * Helpers for reversing scope ranges to gather terminal rewrite information |
| 206 | */ |
| 207 | type ScopeTraversalContext = { |
| 208 | // cache allocated fallthroughs for start/end scope terminal pairs |
| 209 | fallthroughs: Map<ScopeId, BlockId>; |
| 210 | rewrites: Array<TerminalRewriteInfo>; |
| 211 | env: Environment; |
| 212 | }; |
| 213 | |
| 214 | function pushStartScopeTerminal( |
| 215 | scope: ReactiveScope, |
| 216 | context: ScopeTraversalContext, |
| 217 | ): void { |
| 218 | const blockId = context.env.nextBlockId; |
| 219 | const fallthroughId = context.env.nextBlockId; |
| 220 | context.rewrites.push({ |
| 221 | kind: 'StartScope', |
| 222 | blockId, |
| 223 | fallthroughId, |
| 224 | instrId: scope.range.start, |
| 225 | scope, |
| 226 | }); |
| 227 | context.fallthroughs.set(scope.id, fallthroughId); |
| 228 | } |
| 229 | |
| 230 | function pushEndScopeTerminal( |
| 231 | scope: ReactiveScope, |
| 232 | context: ScopeTraversalContext, |
| 233 | ): void { |
| 234 | const fallthroughId = context.fallthroughs.get(scope.id); |
| 235 | CompilerError.invariant(fallthroughId != null, { |
| 236 | reason: 'Expected scope to exist', |
| 237 | loc: GeneratedSource, |
| 238 | }); |
| 239 | context.rewrites.push({ |
| 240 | kind: 'EndScope', |
| 241 | fallthroughId, |
| 242 | instrId: scope.range.end, |
| 243 | }); |
| 244 | } |
| 245 | |
| 246 | type RewriteContext = { |
| 247 | source: BasicBlock; |
| 248 | instrSliceIdx: number; |
| 249 | nextPreds: Set<BlockId>; |
| 250 | nextBlockId: BlockId; |
| 251 | rewrites: Array<BasicBlock>; |
| 252 | }; |
| 253 | |
| 254 | /** |
| 255 | * Create a block rewrite by slicing a set of instructions from source. |
| 256 | * Since scope start-ends always end with a GOTO to the next instruction |
| 257 | * from the source block, we directly connect rewritten blocks using state |
| 258 | * from `context`. |
| 259 | * |
| 260 | * Source: |
| 261 | * bb0: |
| 262 | * instr1, instr2, instr3, instr4, [[ original terminal ]] |
| 263 | * Rewritten: |
| 264 | * bb0: |
| 265 | * instr1, [[ scope start block=bb1]] |
| 266 | * bb1: |
| 267 | * instr2, instr3, [[ scope end goto=bb2 ]] |
| 268 | * bb2: |
| 269 | * instr4, [[ original terminal ]] |
| 270 | */ |
| 271 | function handleRewrite( |
| 272 | terminalInfo: TerminalRewriteInfo, |
| 273 | idx: number, |
| 274 | context: RewriteContext, |
| 275 | ): void { |
| 276 | // TODO make consistent instruction IDs instead of reusing |
| 277 | const terminal: ReactiveScopeTerminal | GotoTerminal = |
| 278 | terminalInfo.kind === 'StartScope' |
| 279 | ? { |
| 280 | kind: 'scope', |
| 281 | fallthrough: terminalInfo.fallthroughId, |
| 282 | block: terminalInfo.blockId, |
| 283 | scope: terminalInfo.scope, |
| 284 | id: terminalInfo.instrId, |
| 285 | loc: GeneratedSource, |
| 286 | } |
| 287 | : { |
| 288 | kind: 'goto', |
| 289 | variant: GotoVariant.Break, |
| 290 | block: terminalInfo.fallthroughId, |
| 291 | id: terminalInfo.instrId, |
| 292 | loc: GeneratedSource, |
| 293 | }; |
| 294 | |
| 295 | const currBlockId = context.nextBlockId; |
| 296 | context.rewrites.push({ |
| 297 | kind: context.source.kind, |
| 298 | id: currBlockId, |
| 299 | instructions: context.source.instructions.slice(context.instrSliceIdx, idx), |
| 300 | preds: context.nextPreds, |
| 301 | // Only the first rewrite should reuse source block phis |
| 302 | phis: context.rewrites.length === 0 ? context.source.phis : new Set(), |
| 303 | terminal, |
| 304 | }); |
| 305 | context.nextPreds = new Set([currBlockId]); |
| 306 | context.nextBlockId = |
| 307 | terminalInfo.kind === 'StartScope' |
| 308 | ? terminalInfo.blockId |
| 309 | : terminalInfo.fallthroughId; |
| 310 | context.instrSliceIdx = idx; |
| 311 | } |