| 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 { |
| 9 | BasicBlock, |
| 10 | BlockId, |
| 11 | Environment, |
| 12 | FunctionExpression, |
| 13 | GeneratedSource, |
| 14 | GotoTerminal, |
| 15 | GotoVariant, |
| 16 | HIRFunction, |
| 17 | IdentifierId, |
| 18 | InstructionKind, |
| 19 | LabelTerminal, |
| 20 | Place, |
| 21 | isStatementBlockKind, |
| 22 | makeInstructionId, |
| 23 | mergeConsecutiveBlocks, |
| 24 | promoteTemporary, |
| 25 | reversePostorderBlocks, |
| 26 | } from '../HIR'; |
| 27 | import { |
| 28 | createTemporaryPlace, |
| 29 | markInstructionIds, |
| 30 | markPredecessors, |
| 31 | } from '../HIR/HIRBuilder'; |
| 32 | import {eachInstructionValueOperand} from '../HIR/visitors'; |
| 33 | import {retainWhere} from '../Utils/utils'; |
| 34 | |
| 35 | /* |
| 36 | * Inlines immediately invoked function expressions (IIFEs) to allow more fine-grained memoization |
| 37 | * of the values they produce. |
| 38 | * |
| 39 | * Example: |
| 40 | * |
| 41 | * ``` |
| 42 | * const x = (() => { |
| 43 | * const x = []; |
| 44 | * x.push(foo()); |
| 45 | * return x; |
| 46 | * })(); |
| 47 | * |
| 48 | * => |
| 49 | * |
| 50 | * bb0: |
| 51 | * // placeholder for the result, all return statements will assign here |
| 52 | * let t0; |
| 53 | * // Label allows using a goto (break) to exit out of the body |
| 54 | * Label block=bb1 fallthrough=bb2 |
| 55 | * bb1: |
| 56 | * // code within the function expression |
| 57 | * const x0 = []; |
| 58 | * x0.push(foo()); |
| 59 | * // return is replaced by assignment to the result variable... |
| 60 | * t0 = x0; |
| 61 | * // ...and a goto to the code after the function expression invocation |
| 62 | * Goto bb2 |
| 63 | * bb2: |
| 64 | * // code after the IIFE call |
| 65 | * const x = t0; |
| 66 | * ``` |
| 67 | * |
| 68 | * The implementation relies on HIR's ability to support labeled blocks: |
| 69 | * - We terminate the basic block just prior to the CallExpression of the IIFE |
| 70 | * with a LabelTerminal whose fallback is the code following the CallExpression. |
| 71 | * Just prior to the terminal we also create a named temporary variable which |
| 72 | * will hold the result. |
| 73 | * - We then inline the contents of the function "in between" (conceptually) those |
| 74 | * two blocks. |
| 75 | * - All return statements in the original function expression are replaced with a |
| 76 | * StoreLocal to the temporary we allocated before plus a Goto to the fallthrough |
| 77 | * block (code following the CallExpression). |
| 78 | * |
| 79 | * Note that if the inliined function has only one return, we avoid the labeled block |
| 80 | * and fully inline the code. The original return is replaced with an assignmen to the |
| 81 | * IIFE's call expression lvalue. |
| 82 | */ |
| 83 | export function inlineImmediatelyInvokedFunctionExpressions( |
| 84 | fn: HIRFunction, |
| 85 | ): void { |
| 86 | // Track all function expressions that are assigned to a temporary |
| 87 | const functions = new Map<IdentifierId, FunctionExpression>(); |
| 88 | // Functions that are inlined |
| 89 | const inlinedFunctions = new Set<IdentifierId>(); |
| 90 | |
| 91 | /* |
| 92 | * Iterate the *existing* blocks from the outer component to find IIFEs |
| 93 | * and inline them. During iteration we will modify `fn` (by inlining the CFG |
| 94 | * of IIFEs) so we explicitly copy references to just the original |
| 95 | * function's blocks first. As blocks are split to make room for IIFE calls, |
| 96 | * the split portions of the blocks will be added to this queue. |
| 97 | */ |
| 98 | const queue = Array.from(fn.body.blocks.values()); |
| 99 | queue: for (const block of queue) { |
| 100 | /* |
| 101 | * We can't handle labels inside expressions yet, so we don't inline IIFEs if they are in an |
| 102 | * expression block. |
| 103 | */ |
| 104 | if (isStatementBlockKind(block.kind)) { |
| 105 | for (let ii = 0; ii < block.instructions.length; ii++) { |
| 106 | const instr = block.instructions[ii]!; |
| 107 | switch (instr.value.kind) { |
| 108 | case 'FunctionExpression': { |
| 109 | if (instr.lvalue.identifier.name === null) { |
| 110 | functions.set(instr.lvalue.identifier.id, instr.value); |
| 111 | } |
| 112 | break; |
| 113 | } |
| 114 | case 'CallExpression': { |
| 115 | if (instr.value.args.length !== 0) { |
| 116 | // We don't support inlining when there are arguments |
| 117 | continue; |
| 118 | } |
| 119 | const body = functions.get(instr.value.callee.identifier.id); |
| 120 | if (body === undefined) { |
| 121 | // Not invoking a local function expression, can't inline |
| 122 | continue; |
| 123 | } |
| 124 | |
| 125 | if ( |
| 126 | body.loweredFunc.func.params.length > 0 || |
| 127 | body.loweredFunc.func.async || |
| 128 | body.loweredFunc.func.generator |
| 129 | ) { |
| 130 | // Can't inline functions with params, or async/generator functions |
| 131 | continue; |
| 132 | } |
| 133 | |
| 134 | // We know this function is used for an IIFE and can prune it later |
| 135 | inlinedFunctions.add(instr.value.callee.identifier.id); |
| 136 | |
| 137 | // Create a new block which will contain code following the IIFE call |
| 138 | const continuationBlockId = fn.env.nextBlockId; |
| 139 | const continuationBlock: BasicBlock = { |
| 140 | id: continuationBlockId, |
| 141 | instructions: block.instructions.slice(ii + 1), |
| 142 | kind: block.kind, |
| 143 | phis: new Set(), |
| 144 | preds: new Set(), |
| 145 | terminal: block.terminal, |
| 146 | }; |
| 147 | fn.body.blocks.set(continuationBlockId, continuationBlock); |
| 148 | |
| 149 | /* |
| 150 | * Trim the original block to contain instructions up to (but not including) |
| 151 | * the IIFE |
| 152 | */ |
| 153 | block.instructions.length = ii; |
| 154 | |
| 155 | if (hasSingleExitReturnTerminal(body.loweredFunc.func)) { |
| 156 | block.terminal = { |
| 157 | kind: 'goto', |
| 158 | block: body.loweredFunc.func.body.entry, |
| 159 | id: block.terminal.id, |
| 160 | loc: block.terminal.loc, |
| 161 | variant: GotoVariant.Break, |
| 162 | } as GotoTerminal; |
| 163 | for (const block of body.loweredFunc.func.body.blocks.values()) { |
| 164 | if (block.terminal.kind === 'return') { |
| 165 | block.instructions.push({ |
| 166 | id: makeInstructionId(0), |
| 167 | loc: block.terminal.loc, |
| 168 | lvalue: instr.lvalue, |
| 169 | value: { |
| 170 | kind: 'LoadLocal', |
| 171 | loc: block.terminal.loc, |
| 172 | place: block.terminal.value, |
| 173 | }, |
| 174 | effects: null, |
| 175 | }); |
| 176 | block.terminal = { |
| 177 | kind: 'goto', |
| 178 | block: continuationBlockId, |
| 179 | id: block.terminal.id, |
| 180 | loc: block.terminal.loc, |
| 181 | variant: GotoVariant.Break, |
| 182 | } as GotoTerminal; |
| 183 | } |
| 184 | } |
| 185 | for (const [id, block] of body.loweredFunc.func.body.blocks) { |
| 186 | block.preds.clear(); |
| 187 | fn.body.blocks.set(id, block); |
| 188 | } |
| 189 | } else { |
| 190 | /* |
| 191 | * To account for multiple returns within the lambda, we treat the lambda |
| 192 | * as if it were a single labeled statement, and replace all returns with gotos |
| 193 | * to the label fallthrough. |
| 194 | */ |
| 195 | const newTerminal: LabelTerminal = { |
| 196 | block: body.loweredFunc.func.body.entry, |
| 197 | id: makeInstructionId(0), |
| 198 | kind: 'label', |
| 199 | fallthrough: continuationBlockId, |
| 200 | loc: block.terminal.loc, |
| 201 | }; |
| 202 | block.terminal = newTerminal; |
| 203 | |
| 204 | // We store the result in the IIFE temporary |
| 205 | const result = instr.lvalue; |
| 206 | |
| 207 | // Declare the IIFE temporary |
| 208 | declareTemporary(fn.env, block, result); |
| 209 | |
| 210 | // Promote the temporary with a name as we require this to persist |
| 211 | if (result.identifier.name == null) { |
| 212 | promoteTemporary(result.identifier); |
| 213 | } |
| 214 | |
| 215 | /* |
| 216 | * Rewrite blocks from the lambda to replace any `return` with a |
| 217 | * store to the result and `goto` the continuation block |
| 218 | */ |
| 219 | for (const [id, block] of body.loweredFunc.func.body.blocks) { |
| 220 | block.preds.clear(); |
| 221 | rewriteBlock(fn.env, block, continuationBlockId, result); |
| 222 | fn.body.blocks.set(id, block); |
| 223 | } |
| 224 | } |
| 225 | |
| 226 | /* |
| 227 | * Ensure we visit the continuation block, since there may have been |
| 228 | * sequential IIFEs that need to be visited. |
| 229 | */ |
| 230 | queue.push(continuationBlock); |
| 231 | continue queue; |
| 232 | } |
| 233 | default: { |
| 234 | for (const place of eachInstructionValueOperand(instr.value)) { |
| 235 | // Any other use of a function expression means it isn't an IIFE |
| 236 | functions.delete(place.identifier.id); |
| 237 | } |
| 238 | } |
| 239 | } |
| 240 | } |
| 241 | } |
| 242 | } |
| 243 | |
| 244 | if (inlinedFunctions.size !== 0) { |
| 245 | // Remove instructions that define lambdas which we inlined |
| 246 | for (const block of fn.body.blocks.values()) { |
| 247 | retainWhere( |
| 248 | block.instructions, |
| 249 | instr => !inlinedFunctions.has(instr.lvalue.identifier.id), |
| 250 | ); |
| 251 | } |
| 252 | |
| 253 | /* |
| 254 | * If terminals have changed then blocks may have become newly unreachable. |
| 255 | * Re-run minification of the graph (incl reordering instruction ids) |
| 256 | */ |
| 257 | reversePostorderBlocks(fn.body); |
| 258 | markInstructionIds(fn.body); |
| 259 | markPredecessors(fn.body); |
| 260 | mergeConsecutiveBlocks(fn); |
| 261 | } |
| 262 | } |
| 263 | |
| 264 | /** |
| 265 | * Returns true if the function has a single exit terminal (throw/return) which is a return |
| 266 | */ |
| 267 | function hasSingleExitReturnTerminal(fn: HIRFunction): boolean { |
| 268 | let hasReturn = false; |
| 269 | let exitCount = 0; |
| 270 | for (const [, block] of fn.body.blocks) { |
| 271 | if (block.terminal.kind === 'return' || block.terminal.kind === 'throw') { |
| 272 | hasReturn ||= block.terminal.kind === 'return'; |
| 273 | exitCount++; |
| 274 | } |
| 275 | } |
| 276 | return exitCount === 1 && hasReturn; |
| 277 | } |
| 278 | |
| 279 | /* |
| 280 | * Rewrites the block so that all `return` terminals are replaced: |
| 281 | * * Add a StoreLocal <returnValue> = <terminal.value> |
| 282 | * * Replace the terminal with a Goto to <returnTarget> |
| 283 | */ |
| 284 | function rewriteBlock( |
| 285 | env: Environment, |
| 286 | block: BasicBlock, |
| 287 | returnTarget: BlockId, |
| 288 | returnValue: Place, |
| 289 | ): void { |
| 290 | const {terminal} = block; |
| 291 | if (terminal.kind !== 'return') { |
| 292 | return; |
| 293 | } |
| 294 | block.instructions.push({ |
| 295 | id: makeInstructionId(0), |
| 296 | loc: terminal.loc, |
| 297 | lvalue: createTemporaryPlace(env, terminal.loc), |
| 298 | value: { |
| 299 | kind: 'StoreLocal', |
| 300 | lvalue: {kind: InstructionKind.Reassign, place: {...returnValue}}, |
| 301 | value: terminal.value, |
| 302 | type: null, |
| 303 | loc: terminal.loc, |
| 304 | }, |
| 305 | effects: null, |
| 306 | }); |
| 307 | block.terminal = { |
| 308 | kind: 'goto', |
| 309 | block: returnTarget, |
| 310 | id: makeInstructionId(0), |
| 311 | variant: GotoVariant.Break, |
| 312 | loc: block.terminal.loc, |
| 313 | }; |
| 314 | } |
| 315 | |
| 316 | function declareTemporary( |
| 317 | env: Environment, |
| 318 | block: BasicBlock, |
| 319 | result: Place, |
| 320 | ): void { |
| 321 | block.instructions.push({ |
| 322 | id: makeInstructionId(0), |
| 323 | loc: GeneratedSource, |
| 324 | lvalue: createTemporaryPlace(env, result.loc), |
| 325 | value: { |
| 326 | kind: 'DeclareLocal', |
| 327 | lvalue: { |
| 328 | place: result, |
| 329 | kind: InstructionKind.Let, |
| 330 | }, |
| 331 | type: null, |
| 332 | loc: result.loc, |
| 333 | }, |
| 334 | effects: null, |
| 335 | }); |
| 336 | } |