| 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 invariant from 'invariant'; |
| 9 | import {Environment} from '../HIR'; |
| 10 | import { |
| 11 | BasicBlock, |
| 12 | GeneratedSource, |
| 13 | HIRFunction, |
| 14 | IdentifierId, |
| 15 | Instruction, |
| 16 | InstructionId, |
| 17 | InstructionKind, |
| 18 | JsxAttribute, |
| 19 | JsxExpression, |
| 20 | LoadGlobal, |
| 21 | makeBlockId, |
| 22 | makeIdentifierName, |
| 23 | makeInstructionId, |
| 24 | ObjectProperty, |
| 25 | Place, |
| 26 | promoteTemporary, |
| 27 | promoteTemporaryJsxTag, |
| 28 | } from '../HIR/HIR'; |
| 29 | import {createTemporaryPlace} from '../HIR/HIRBuilder'; |
| 30 | import {printIdentifier} from '../HIR/PrintHIR'; |
| 31 | import {deadCodeElimination} from './DeadCodeElimination'; |
| 32 | import {assertExhaustive} from '../Utils/utils'; |
| 33 | |
| 34 | export function outlineJSX(fn: HIRFunction): void { |
| 35 | const outlinedFns: Array<HIRFunction> = []; |
| 36 | outlineJsxImpl(fn, outlinedFns); |
| 37 | |
| 38 | for (const outlinedFn of outlinedFns) { |
| 39 | fn.env.outlineFunction(outlinedFn, 'Component'); |
| 40 | } |
| 41 | } |
| 42 | |
| 43 | type JsxInstruction = Instruction & {value: JsxExpression}; |
| 44 | type LoadGlobalInstruction = Instruction & {value: LoadGlobal}; |
| 45 | type LoadGlobalMap = Map<IdentifierId, LoadGlobalInstruction>; |
| 46 | |
| 47 | type State = { |
| 48 | jsx: Array<JsxInstruction>; |
| 49 | children: Set<IdentifierId>; |
| 50 | }; |
| 51 | |
| 52 | function outlineJsxImpl( |
| 53 | fn: HIRFunction, |
| 54 | outlinedFns: Array<HIRFunction>, |
| 55 | ): void { |
| 56 | const globals: LoadGlobalMap = new Map(); |
| 57 | |
| 58 | function processAndOutlineJSX( |
| 59 | state: State, |
| 60 | rewriteInstr: Map<InstructionId, Array<Instruction>>, |
| 61 | ): void { |
| 62 | if (state.jsx.length <= 1) { |
| 63 | return; |
| 64 | } |
| 65 | const result = process( |
| 66 | fn, |
| 67 | [...state.jsx].sort((a, b) => a.id - b.id), |
| 68 | globals, |
| 69 | ); |
| 70 | if (result) { |
| 71 | outlinedFns.push(result.fn); |
| 72 | rewriteInstr.set(state.jsx.at(0)!.id, result.instrs); |
| 73 | } |
| 74 | } |
| 75 | |
| 76 | for (const [, block] of fn.body.blocks) { |
| 77 | const rewriteInstr = new Map(); |
| 78 | let state: State = { |
| 79 | jsx: [], |
| 80 | children: new Set(), |
| 81 | }; |
| 82 | |
| 83 | for (let i = block.instructions.length - 1; i >= 0; i--) { |
| 84 | const instr = block.instructions[i]; |
| 85 | const {value, lvalue} = instr; |
| 86 | switch (value.kind) { |
| 87 | case 'LoadGlobal': { |
| 88 | globals.set(lvalue.identifier.id, instr as LoadGlobalInstruction); |
| 89 | break; |
| 90 | } |
| 91 | case 'FunctionExpression': { |
| 92 | outlineJsxImpl(value.loweredFunc.func, outlinedFns); |
| 93 | break; |
| 94 | } |
| 95 | |
| 96 | case 'JsxExpression': { |
| 97 | if (!state.children.has(lvalue.identifier.id)) { |
| 98 | processAndOutlineJSX(state, rewriteInstr); |
| 99 | |
| 100 | state = { |
| 101 | jsx: [], |
| 102 | children: new Set(), |
| 103 | }; |
| 104 | } |
| 105 | state.jsx.push(instr as JsxInstruction); |
| 106 | if (value.children) { |
| 107 | for (const child of value.children) { |
| 108 | state.children.add(child.identifier.id); |
| 109 | } |
| 110 | } |
| 111 | break; |
| 112 | } |
| 113 | case 'ArrayExpression': |
| 114 | case 'Await': |
| 115 | case 'BinaryExpression': |
| 116 | case 'CallExpression': |
| 117 | case 'ComputedDelete': |
| 118 | case 'ComputedLoad': |
| 119 | case 'ComputedStore': |
| 120 | case 'Debugger': |
| 121 | case 'DeclareContext': |
| 122 | case 'DeclareLocal': |
| 123 | case 'Destructure': |
| 124 | case 'FinishMemoize': |
| 125 | case 'GetIterator': |
| 126 | case 'IteratorNext': |
| 127 | case 'JSXText': |
| 128 | case 'JsxFragment': |
| 129 | case 'LoadContext': |
| 130 | case 'LoadLocal': |
| 131 | case 'MetaProperty': |
| 132 | case 'MethodCall': |
| 133 | case 'NewExpression': |
| 134 | case 'NextPropertyOf': |
| 135 | case 'ObjectExpression': |
| 136 | case 'ObjectMethod': |
| 137 | case 'PostfixUpdate': |
| 138 | case 'PrefixUpdate': |
| 139 | case 'Primitive': |
| 140 | case 'PropertyDelete': |
| 141 | case 'PropertyLoad': |
| 142 | case 'PropertyStore': |
| 143 | case 'RegExpLiteral': |
| 144 | case 'StartMemoize': |
| 145 | case 'StoreContext': |
| 146 | case 'StoreGlobal': |
| 147 | case 'StoreLocal': |
| 148 | case 'TaggedTemplateExpression': |
| 149 | case 'TemplateLiteral': |
| 150 | case 'TypeCastExpression': |
| 151 | case 'UnsupportedNode': |
| 152 | case 'UnaryExpression': { |
| 153 | break; |
| 154 | } |
| 155 | default: { |
| 156 | assertExhaustive(value, `Unexpected instruction: ${value}`); |
| 157 | } |
| 158 | } |
| 159 | } |
| 160 | processAndOutlineJSX(state, rewriteInstr); |
| 161 | |
| 162 | if (rewriteInstr.size > 0) { |
| 163 | const newInstrs = []; |
| 164 | for (let i = 0; i < block.instructions.length; i++) { |
| 165 | // InstructionId's are one-indexed, so add one to account for them. |
| 166 | const id = i + 1; |
| 167 | if (rewriteInstr.has(id)) { |
| 168 | const instrs = rewriteInstr.get(id); |
| 169 | newInstrs.push(...instrs); |
| 170 | } else { |
| 171 | newInstrs.push(block.instructions[i]); |
| 172 | } |
| 173 | } |
| 174 | block.instructions = newInstrs; |
| 175 | } |
| 176 | deadCodeElimination(fn); |
| 177 | } |
| 178 | } |
| 179 | |
| 180 | type OutlinedResult = { |
| 181 | instrs: Array<Instruction>; |
| 182 | fn: HIRFunction; |
| 183 | }; |
| 184 | |
| 185 | function process( |
| 186 | fn: HIRFunction, |
| 187 | jsx: Array<JsxInstruction>, |
| 188 | globals: LoadGlobalMap, |
| 189 | ): OutlinedResult | null { |
| 190 | /** |
| 191 | * In the future, add a check for backedge to outline jsx inside loops in a |
| 192 | * top level component. For now, only outline jsx in callbacks. |
| 193 | */ |
| 194 | if (fn.fnType === 'Component') { |
| 195 | return null; |
| 196 | } |
| 197 | |
| 198 | const props = collectProps(fn.env, jsx); |
| 199 | if (!props) return null; |
| 200 | |
| 201 | const outlinedTag = fn.env.generateGloballyUniqueIdentifierName(null).value; |
| 202 | const newInstrs = emitOutlinedJsx(fn.env, jsx, props, outlinedTag); |
| 203 | if (!newInstrs) return null; |
| 204 | |
| 205 | const outlinedFn = emitOutlinedFn(fn.env, jsx, props, globals); |
| 206 | if (!outlinedFn) return null; |
| 207 | outlinedFn.id = outlinedTag; |
| 208 | |
| 209 | return {instrs: newInstrs, fn: outlinedFn}; |
| 210 | } |
| 211 | |
| 212 | type OutlinedJsxAttribute = { |
| 213 | originalName: string; |
| 214 | newName: string; |
| 215 | place: Place; |
| 216 | }; |
| 217 | |
| 218 | function collectProps( |
| 219 | env: Environment, |
| 220 | instructions: Array<JsxInstruction>, |
| 221 | ): Array<OutlinedJsxAttribute> | null { |
| 222 | let id = 1; |
| 223 | |
| 224 | function generateName(oldName: string): string { |
| 225 | let newName = oldName; |
| 226 | while (seen.has(newName)) { |
| 227 | newName = `${oldName}${id++}`; |
| 228 | } |
| 229 | seen.add(newName); |
| 230 | env.programContext.addNewReference(newName); |
| 231 | return newName; |
| 232 | } |
| 233 | |
| 234 | const attributes: Array<OutlinedJsxAttribute> = []; |
| 235 | const jsxIds = new Set(instructions.map(i => i.lvalue.identifier.id)); |
| 236 | const seen: Set<string> = new Set(); |
| 237 | |
| 238 | for (const instr of instructions) { |
| 239 | const {value} = instr; |
| 240 | |
| 241 | for (const at of value.props) { |
| 242 | if (at.kind === 'JsxSpreadAttribute') { |
| 243 | return null; |
| 244 | } |
| 245 | |
| 246 | if (at.kind === 'JsxAttribute') { |
| 247 | const newName = generateName(at.name); |
| 248 | attributes.push({ |
| 249 | originalName: at.name, |
| 250 | newName, |
| 251 | place: at.place, |
| 252 | }); |
| 253 | } |
| 254 | } |
| 255 | |
| 256 | if (value.children) { |
| 257 | for (const child of value.children) { |
| 258 | if (jsxIds.has(child.identifier.id)) { |
| 259 | continue; |
| 260 | } |
| 261 | |
| 262 | promoteTemporary(child.identifier); |
| 263 | const newName = generateName('t'); |
| 264 | attributes.push({ |
| 265 | originalName: child.identifier.name!.value, |
| 266 | newName: newName, |
| 267 | place: child, |
| 268 | }); |
| 269 | } |
| 270 | } |
| 271 | } |
| 272 | return attributes; |
| 273 | } |
| 274 | |
| 275 | function emitOutlinedJsx( |
| 276 | env: Environment, |
| 277 | instructions: Array<Instruction>, |
| 278 | outlinedProps: Array<OutlinedJsxAttribute>, |
| 279 | outlinedTag: string, |
| 280 | ): Array<Instruction> { |
| 281 | const props: Array<JsxAttribute> = outlinedProps.map(p => ({ |
| 282 | kind: 'JsxAttribute', |
| 283 | name: p.newName, |
| 284 | place: p.place, |
| 285 | })); |
| 286 | |
| 287 | const loadJsx: Instruction = { |
| 288 | id: makeInstructionId(0), |
| 289 | loc: GeneratedSource, |
| 290 | lvalue: createTemporaryPlace(env, GeneratedSource), |
| 291 | value: { |
| 292 | kind: 'LoadGlobal', |
| 293 | binding: { |
| 294 | kind: 'ModuleLocal', |
| 295 | name: outlinedTag, |
| 296 | }, |
| 297 | loc: GeneratedSource, |
| 298 | }, |
| 299 | effects: null, |
| 300 | }; |
| 301 | promoteTemporaryJsxTag(loadJsx.lvalue.identifier); |
| 302 | const jsxExpr: Instruction = { |
| 303 | id: makeInstructionId(0), |
| 304 | loc: GeneratedSource, |
| 305 | lvalue: instructions.at(-1)!.lvalue, |
| 306 | value: { |
| 307 | kind: 'JsxExpression', |
| 308 | tag: {...loadJsx.lvalue}, |
| 309 | props, |
| 310 | children: null, |
| 311 | loc: GeneratedSource, |
| 312 | openingLoc: GeneratedSource, |
| 313 | closingLoc: GeneratedSource, |
| 314 | }, |
| 315 | effects: null, |
| 316 | }; |
| 317 | |
| 318 | return [loadJsx, jsxExpr]; |
| 319 | } |
| 320 | |
| 321 | function emitOutlinedFn( |
| 322 | env: Environment, |
| 323 | jsx: Array<JsxInstruction>, |
| 324 | oldProps: Array<OutlinedJsxAttribute>, |
| 325 | globals: LoadGlobalMap, |
| 326 | ): HIRFunction | null { |
| 327 | const instructions: Array<Instruction> = []; |
| 328 | const oldToNewProps = createOldToNewPropsMapping(env, oldProps); |
| 329 | |
| 330 | const propsObj: Place = createTemporaryPlace(env, GeneratedSource); |
| 331 | promoteTemporary(propsObj.identifier); |
| 332 | |
| 333 | const destructurePropsInstr = emitDestructureProps( |
| 334 | env, |
| 335 | propsObj, |
| 336 | oldToNewProps, |
| 337 | ); |
| 338 | instructions.push(destructurePropsInstr); |
| 339 | |
| 340 | const updatedJsxInstructions = emitUpdatedJsx(jsx, oldToNewProps); |
| 341 | const loadGlobalInstrs = emitLoadGlobals(jsx, globals); |
| 342 | if (!loadGlobalInstrs) { |
| 343 | return null; |
| 344 | } |
| 345 | instructions.push(...loadGlobalInstrs); |
| 346 | instructions.push(...updatedJsxInstructions); |
| 347 | |
| 348 | const block: BasicBlock = { |
| 349 | kind: 'block', |
| 350 | id: makeBlockId(0), |
| 351 | instructions, |
| 352 | terminal: { |
| 353 | id: makeInstructionId(0), |
| 354 | kind: 'return', |
| 355 | returnVariant: 'Explicit', |
| 356 | loc: GeneratedSource, |
| 357 | value: instructions.at(-1)!.lvalue, |
| 358 | effects: null, |
| 359 | }, |
| 360 | preds: new Set(), |
| 361 | phis: new Set(), |
| 362 | }; |
| 363 | |
| 364 | const fn: HIRFunction = { |
| 365 | loc: GeneratedSource, |
| 366 | id: null, |
| 367 | nameHint: null, |
| 368 | fnType: 'Other', |
| 369 | env, |
| 370 | params: [propsObj], |
| 371 | returnTypeAnnotation: null, |
| 372 | returns: createTemporaryPlace(env, GeneratedSource), |
| 373 | context: [], |
| 374 | body: { |
| 375 | entry: block.id, |
| 376 | blocks: new Map([[block.id, block]]), |
| 377 | }, |
| 378 | generator: false, |
| 379 | async: false, |
| 380 | directives: [], |
| 381 | aliasingEffects: [], |
| 382 | }; |
| 383 | return fn; |
| 384 | } |
| 385 | |
| 386 | function emitLoadGlobals( |
| 387 | jsx: Array<JsxInstruction>, |
| 388 | globals: LoadGlobalMap, |
| 389 | ): Array<Instruction> | null { |
| 390 | const instructions: Array<Instruction> = []; |
| 391 | for (const {value} of jsx) { |
| 392 | // Add load globals instructions for jsx tags |
| 393 | if (value.tag.kind === 'Identifier') { |
| 394 | const loadGlobalInstr = globals.get(value.tag.identifier.id); |
| 395 | if (!loadGlobalInstr) { |
| 396 | return null; |
| 397 | } |
| 398 | instructions.push(loadGlobalInstr); |
| 399 | } |
| 400 | } |
| 401 | |
| 402 | return instructions; |
| 403 | } |
| 404 | |
| 405 | function emitUpdatedJsx( |
| 406 | jsx: Array<JsxInstruction>, |
| 407 | oldToNewProps: Map<IdentifierId, OutlinedJsxAttribute>, |
| 408 | ): Array<JsxInstruction> { |
| 409 | const newInstrs: Array<JsxInstruction> = []; |
| 410 | const jsxIds = new Set(jsx.map(i => i.lvalue.identifier.id)); |
| 411 | |
| 412 | for (const instr of jsx) { |
| 413 | const {value} = instr; |
| 414 | const newProps: Array<JsxAttribute> = []; |
| 415 | // Update old props references to use the newly destructured props param |
| 416 | for (const prop of value.props) { |
| 417 | invariant( |
| 418 | prop.kind === 'JsxAttribute', |
| 419 | `Expected only attributes but found ${prop.kind}`, |
| 420 | ); |
| 421 | if (prop.name === 'key') { |
| 422 | continue; |
| 423 | } |
| 424 | const newProp = oldToNewProps.get(prop.place.identifier.id); |
| 425 | invariant( |
| 426 | newProp !== undefined, |
| 427 | `Expected a new property for ${printIdentifier(prop.place.identifier)}`, |
| 428 | ); |
| 429 | newProps.push({ |
| 430 | kind: 'JsxAttribute', |
| 431 | name: newProp.originalName, |
| 432 | place: newProp.place, |
| 433 | }); |
| 434 | } |
| 435 | |
| 436 | let newChildren: Array<Place> | null = null; |
| 437 | if (value.children) { |
| 438 | newChildren = []; |
| 439 | for (const child of value.children) { |
| 440 | if (jsxIds.has(child.identifier.id)) { |
| 441 | newChildren.push({...child}); |
| 442 | continue; |
| 443 | } |
| 444 | |
| 445 | const newChild = oldToNewProps.get(child.identifier.id); |
| 446 | invariant( |
| 447 | newChild !== undefined, |
| 448 | `Expected a new prop for ${printIdentifier(child.identifier)}`, |
| 449 | ); |
| 450 | newChildren.push({...newChild.place}); |
| 451 | } |
| 452 | } |
| 453 | |
| 454 | newInstrs.push({ |
| 455 | ...instr, |
| 456 | value: { |
| 457 | ...value, |
| 458 | props: newProps, |
| 459 | children: newChildren, |
| 460 | }, |
| 461 | }); |
| 462 | } |
| 463 | |
| 464 | return newInstrs; |
| 465 | } |
| 466 | |
| 467 | function createOldToNewPropsMapping( |
| 468 | env: Environment, |
| 469 | oldProps: Array<OutlinedJsxAttribute>, |
| 470 | ): Map<IdentifierId, OutlinedJsxAttribute> { |
| 471 | const oldToNewProps = new Map(); |
| 472 | |
| 473 | for (const oldProp of oldProps) { |
| 474 | // Do not read key prop in the outlined component |
| 475 | if (oldProp.originalName === 'key') { |
| 476 | continue; |
| 477 | } |
| 478 | |
| 479 | const newProp: OutlinedJsxAttribute = { |
| 480 | ...oldProp, |
| 481 | place: createTemporaryPlace(env, GeneratedSource), |
| 482 | }; |
| 483 | newProp.place.identifier.name = makeIdentifierName(oldProp.newName); |
| 484 | oldToNewProps.set(oldProp.place.identifier.id, newProp); |
| 485 | } |
| 486 | |
| 487 | return oldToNewProps; |
| 488 | } |
| 489 | |
| 490 | function emitDestructureProps( |
| 491 | env: Environment, |
| 492 | propsObj: Place, |
| 493 | oldToNewProps: Map<IdentifierId, OutlinedJsxAttribute>, |
| 494 | ): Instruction { |
| 495 | const properties: Array<ObjectProperty> = []; |
| 496 | for (const [_, prop] of oldToNewProps) { |
| 497 | properties.push({ |
| 498 | kind: 'ObjectProperty', |
| 499 | key: { |
| 500 | kind: 'string', |
| 501 | name: prop.newName, |
| 502 | }, |
| 503 | type: 'property', |
| 504 | place: prop.place, |
| 505 | }); |
| 506 | } |
| 507 | |
| 508 | const destructurePropsInstr: Instruction = { |
| 509 | id: makeInstructionId(0), |
| 510 | lvalue: createTemporaryPlace(env, GeneratedSource), |
| 511 | loc: GeneratedSource, |
| 512 | value: { |
| 513 | kind: 'Destructure', |
| 514 | lvalue: { |
| 515 | pattern: { |
| 516 | kind: 'ObjectPattern', |
| 517 | properties, |
| 518 | loc: GeneratedSource, |
| 519 | }, |
| 520 | kind: InstructionKind.Let, |
| 521 | }, |
| 522 | loc: GeneratedSource, |
| 523 | value: propsObj, |
| 524 | }, |
| 525 | effects: null, |
| 526 | }; |
| 527 | return destructurePropsInstr; |
| 528 | } |