| 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 {Environment} from '../HIR/Environment'; |
| 10 | import { |
| 11 | BasicBlock, |
| 12 | BlockId, |
| 13 | GeneratedSource, |
| 14 | HIRFunction, |
| 15 | Identifier, |
| 16 | IdentifierId, |
| 17 | makeInstructionId, |
| 18 | makeType, |
| 19 | Phi, |
| 20 | Place, |
| 21 | } from '../HIR/HIR'; |
| 22 | import {printIdentifier, printPlace} from '../HIR/PrintHIR'; |
| 23 | import { |
| 24 | eachTerminalSuccessor, |
| 25 | mapInstructionLValues, |
| 26 | mapInstructionOperands, |
| 27 | mapTerminalOperands, |
| 28 | } from '../HIR/visitors'; |
| 29 | |
| 30 | type IncompletePhi = { |
| 31 | oldPlace: Place; |
| 32 | newPlace: Place; |
| 33 | }; |
| 34 | |
| 35 | type State = { |
| 36 | defs: Map<Identifier, Identifier>; |
| 37 | incompletePhis: Array<IncompletePhi>; |
| 38 | }; |
| 39 | |
| 40 | class SSABuilder { |
| 41 | #states: Map<BasicBlock, State> = new Map(); |
| 42 | #current: BasicBlock | null = null; |
| 43 | unsealedPreds: Map<BasicBlock, number> = new Map(); |
| 44 | #blocks: Map<BlockId, BasicBlock>; |
| 45 | #env: Environment; |
| 46 | #unknown: Set<Identifier> = new Set(); |
| 47 | #context: Set<Identifier> = new Set(); |
| 48 | |
| 49 | constructor(env: Environment, blocks: ReadonlyMap<BlockId, BasicBlock>) { |
| 50 | this.#blocks = new Map(blocks); |
| 51 | this.#env = env; |
| 52 | } |
| 53 | |
| 54 | get nextSsaId(): IdentifierId { |
| 55 | return this.#env.nextIdentifierId; |
| 56 | } |
| 57 | |
| 58 | defineFunction(func: HIRFunction): void { |
| 59 | for (const [id, block] of func.body.blocks) { |
| 60 | this.#blocks.set(id, block); |
| 61 | } |
| 62 | } |
| 63 | |
| 64 | enter(fn: () => void): void { |
| 65 | const current = this.#current; |
| 66 | fn(); |
| 67 | this.#current = current; |
| 68 | } |
| 69 | |
| 70 | state(): State { |
| 71 | CompilerError.invariant(this.#current !== null, { |
| 72 | reason: 'we need to be in a block to access state!', |
| 73 | loc: GeneratedSource, |
| 74 | }); |
| 75 | return this.#states.get(this.#current)!; |
| 76 | } |
| 77 | |
| 78 | makeId(oldId: Identifier): Identifier { |
| 79 | return { |
| 80 | id: this.nextSsaId, |
| 81 | declarationId: oldId.declarationId, |
| 82 | name: oldId.name, |
| 83 | mutableRange: { |
| 84 | start: makeInstructionId(0), |
| 85 | end: makeInstructionId(0), |
| 86 | }, |
| 87 | scope: null, // reset along w the mutable range |
| 88 | type: makeType(), |
| 89 | loc: oldId.loc, |
| 90 | }; |
| 91 | } |
| 92 | |
| 93 | defineContext(oldPlace: Place): Place { |
| 94 | const newPlace = this.definePlace(oldPlace); |
| 95 | this.#context.add(oldPlace.identifier); |
| 96 | return newPlace; |
| 97 | } |
| 98 | |
| 99 | /** |
| 100 | * A function's context places capture a *binding*, not a value: the |
| 101 | * variable is only read when the function is later called, so a context |
| 102 | * place may reference a binding that is declared after the function |
| 103 | * expression itself (eg `const colgroup = useMemo(() => <colgroup>...)`, |
| 104 | * where the JSX tag name resolves to the variable being assigned). Unmark |
| 105 | * such identifiers so the later declaration doesn't error; if the function |
| 106 | * body actually *reads* the variable before it is defined, visiting the |
| 107 | * body re-marks it and the hoisting bailout in definePlace still applies. |
| 108 | */ |
| 109 | unmarkUnknown(place: Place): void { |
| 110 | this.#unknown.delete(place.identifier); |
| 111 | } |
| 112 | |
| 113 | definePlace(oldPlace: Place): Place { |
| 114 | const oldId = oldPlace.identifier; |
| 115 | if (this.#unknown.has(oldId)) { |
| 116 | CompilerError.throwTodo({ |
| 117 | reason: `[hoisting] EnterSSA: Expected identifier to be defined before being used`, |
| 118 | description: `Identifier ${printIdentifier(oldId)} is undefined`, |
| 119 | loc: oldPlace.loc, |
| 120 | suggestions: null, |
| 121 | }); |
| 122 | } |
| 123 | |
| 124 | // Do not redefine context references. |
| 125 | if (this.#context.has(oldId)) { |
| 126 | return this.getPlace(oldPlace); |
| 127 | } |
| 128 | |
| 129 | const newId = this.makeId(oldId); |
| 130 | this.state().defs.set(oldId, newId); |
| 131 | return { |
| 132 | ...oldPlace, |
| 133 | identifier: newId, |
| 134 | }; |
| 135 | } |
| 136 | |
| 137 | getPlace(oldPlace: Place): Place { |
| 138 | const newId = this.getIdAt(oldPlace, this.#current!.id); |
| 139 | return { |
| 140 | ...oldPlace, |
| 141 | identifier: newId, |
| 142 | }; |
| 143 | } |
| 144 | |
| 145 | getIdAt(oldPlace: Place, blockId: BlockId): Identifier { |
| 146 | // check if Place is defined locally |
| 147 | const block = this.#blocks.get(blockId)!; |
| 148 | const state = this.#states.get(block)!; |
| 149 | |
| 150 | if (state.defs.has(oldPlace.identifier)) { |
| 151 | return state.defs.get(oldPlace.identifier)!; |
| 152 | } |
| 153 | |
| 154 | if (block.preds.size == 0) { |
| 155 | /* |
| 156 | * We're at the entry block and haven't found our defintion yet. |
| 157 | * console.log( |
| 158 | * `Unable to find "${printPlace( |
| 159 | * oldPlace |
| 160 | * )}" in bb${blockId}, assuming it's a global` |
| 161 | * ); |
| 162 | */ |
| 163 | this.#unknown.add(oldPlace.identifier); |
| 164 | return oldPlace.identifier; |
| 165 | } |
| 166 | |
| 167 | if (this.unsealedPreds.get(block)! > 0) { |
| 168 | /* |
| 169 | * We haven't visited all our predecessors, let's place an incomplete phi |
| 170 | * for now. |
| 171 | */ |
| 172 | const newId = this.makeId(oldPlace.identifier); |
| 173 | state.incompletePhis.push({ |
| 174 | oldPlace, |
| 175 | newPlace: {...oldPlace, identifier: newId}, |
| 176 | }); |
| 177 | state.defs.set(oldPlace.identifier, newId); |
| 178 | return newId; |
| 179 | } |
| 180 | |
| 181 | // Only one predecessor, let's check there |
| 182 | if (block.preds.size == 1) { |
| 183 | const [pred] = block.preds; |
| 184 | const newId = this.getIdAt(oldPlace, pred); |
| 185 | state.defs.set(oldPlace.identifier, newId); |
| 186 | return newId; |
| 187 | } |
| 188 | |
| 189 | // There are multiple predecessors, we may need a phi. |
| 190 | const newId = this.makeId(oldPlace.identifier); |
| 191 | /* |
| 192 | * Adding a phi may loop back to our block if there is a loop in the CFG. We |
| 193 | * update our defs before adding the phi to terminate the recursion rather than |
| 194 | * looping infinitely. |
| 195 | */ |
| 196 | state.defs.set(oldPlace.identifier, newId); |
| 197 | return this.addPhi(block, oldPlace, {...oldPlace, identifier: newId}); |
| 198 | } |
| 199 | |
| 200 | addPhi(block: BasicBlock, oldPlace: Place, newPlace: Place): Identifier { |
| 201 | const predDefs: Map<BlockId, Place> = new Map(); |
| 202 | for (const predBlockId of block.preds) { |
| 203 | const predId = this.getIdAt(oldPlace, predBlockId); |
| 204 | predDefs.set(predBlockId, {...oldPlace, identifier: predId}); |
| 205 | } |
| 206 | |
| 207 | const phi: Phi = { |
| 208 | kind: 'Phi', |
| 209 | place: newPlace, |
| 210 | operands: predDefs, |
| 211 | }; |
| 212 | |
| 213 | block.phis.add(phi); |
| 214 | return newPlace.identifier; |
| 215 | } |
| 216 | |
| 217 | fixIncompletePhis(block: BasicBlock): void { |
| 218 | const state = this.#states.get(block)!; |
| 219 | for (const phi of state.incompletePhis) { |
| 220 | this.addPhi(block, phi.oldPlace, phi.newPlace); |
| 221 | } |
| 222 | } |
| 223 | |
| 224 | startBlock(block: BasicBlock): void { |
| 225 | this.#current = block; |
| 226 | this.#states.set(block, { |
| 227 | defs: new Map(), |
| 228 | incompletePhis: [], |
| 229 | }); |
| 230 | } |
| 231 | |
| 232 | print(): void { |
| 233 | const text: Array<string> = []; |
| 234 | for (const [block, state] of this.#states) { |
| 235 | text.push(`bb${block.id}:`); |
| 236 | for (const [oldId, newId] of state.defs) { |
| 237 | text.push(` \$${printIdentifier(oldId)}: \$${printIdentifier(newId)}`); |
| 238 | } |
| 239 | |
| 240 | for (const incompletePhi of state.incompletePhis) { |
| 241 | text.push( |
| 242 | ` iphi \$${printPlace( |
| 243 | incompletePhi.newPlace, |
| 244 | )} = \$${printPlace(incompletePhi.oldPlace)}`, |
| 245 | ); |
| 246 | } |
| 247 | } |
| 248 | |
| 249 | text.push(`current block: bb${this.#current?.id}`); |
| 250 | console.log(text.join('\n')); |
| 251 | } |
| 252 | } |
| 253 | |
| 254 | export default function enterSSA(func: HIRFunction): void { |
| 255 | const builder = new SSABuilder(func.env, func.body.blocks); |
| 256 | enterSSAImpl(func, builder, func.body.entry); |
| 257 | } |
| 258 | |
| 259 | function enterSSAImpl( |
| 260 | func: HIRFunction, |
| 261 | builder: SSABuilder, |
| 262 | rootEntry: BlockId, |
| 263 | ): void { |
| 264 | const visitedBlocks: Set<BasicBlock> = new Set(); |
| 265 | for (const [blockId, block] of func.body.blocks) { |
| 266 | CompilerError.invariant(!visitedBlocks.has(block), { |
| 267 | reason: `found a cycle! visiting bb${block.id} again`, |
| 268 | loc: GeneratedSource, |
| 269 | }); |
| 270 | |
| 271 | visitedBlocks.add(block); |
| 272 | |
| 273 | builder.startBlock(block); |
| 274 | |
| 275 | if (blockId === rootEntry) { |
| 276 | // NOTE: func.context should be empty for the root function |
| 277 | CompilerError.invariant(func.context.length === 0, { |
| 278 | reason: `Expected function context to be empty for outer function declarations`, |
| 279 | loc: func.loc, |
| 280 | }); |
| 281 | func.params = func.params.map(param => { |
| 282 | if (param.kind === 'Identifier') { |
| 283 | return builder.definePlace(param); |
| 284 | } else { |
| 285 | return { |
| 286 | kind: 'Spread', |
| 287 | place: builder.definePlace(param.place), |
| 288 | }; |
| 289 | } |
| 290 | }); |
| 291 | } |
| 292 | |
| 293 | for (const instr of block.instructions) { |
| 294 | mapInstructionOperands(instr, place => builder.getPlace(place)); |
| 295 | mapInstructionLValues(instr, lvalue => builder.definePlace(lvalue)); |
| 296 | |
| 297 | if ( |
| 298 | instr.value.kind === 'FunctionExpression' || |
| 299 | instr.value.kind === 'ObjectMethod' |
| 300 | ) { |
| 301 | const loweredFunc = instr.value.loweredFunc.func; |
| 302 | for (const place of loweredFunc.context) { |
| 303 | builder.unmarkUnknown(place); |
| 304 | } |
| 305 | const entry = loweredFunc.body.blocks.get(loweredFunc.body.entry)!; |
| 306 | CompilerError.invariant(entry.preds.size === 0, { |
| 307 | reason: |
| 308 | 'Expected function expression entry block to have zero predecessors', |
| 309 | loc: GeneratedSource, |
| 310 | }); |
| 311 | entry.preds.add(blockId); |
| 312 | builder.defineFunction(loweredFunc); |
| 313 | builder.enter(() => { |
| 314 | loweredFunc.params = loweredFunc.params.map(param => { |
| 315 | if (param.kind === 'Identifier') { |
| 316 | return builder.definePlace(param); |
| 317 | } else { |
| 318 | return { |
| 319 | kind: 'Spread', |
| 320 | place: builder.definePlace(param.place), |
| 321 | }; |
| 322 | } |
| 323 | }); |
| 324 | enterSSAImpl(loweredFunc, builder, rootEntry); |
| 325 | }); |
| 326 | entry.preds.clear(); |
| 327 | } |
| 328 | } |
| 329 | |
| 330 | mapTerminalOperands(block.terminal, place => builder.getPlace(place)); |
| 331 | for (const outputId of eachTerminalSuccessor(block.terminal)) { |
| 332 | const output = func.body.blocks.get(outputId)!; |
| 333 | let count; |
| 334 | if (builder.unsealedPreds.has(output)) { |
| 335 | count = builder.unsealedPreds.get(output)! - 1; |
| 336 | } else { |
| 337 | count = output.preds.size - 1; |
| 338 | } |
| 339 | builder.unsealedPreds.set(output, count); |
| 340 | |
| 341 | if (count === 0 && visitedBlocks.has(output)) { |
| 342 | builder.fixIncompletePhis(output); |
| 343 | } |
| 344 | } |
| 345 | } |
| 346 | } |