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