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, CompilerErrorDetailOptions, ErrorSeverity} from '..';
9
+import {
10
+ CallExpression,
11
+ Effect,
12
+ Environment,
13
+ FunctionExpression,
14
+ GeneratedSource,
15
+ HIRFunction,
16
+ Identifier,
17
+ IdentifierId,
18
+ Instruction,
19
+ InstructionId,
20
+ InstructionKind,
21
+ InstructionValue,
22
+ isUseEffectHookType,
23
+ LoadLocal,
24
+ makeInstructionId,
25
+ Place,
26
+ promoteTemporary,
27
+} from '../HIR';
28
+import {createTemporaryPlace, markInstructionIds} from '../HIR/HIRBuilder';
29
+import {getOrInsertWith} from '../Utils/utils';
30
+import {BuiltInFireId, DefaultNonmutatingHook} from '../HIR/ObjectShape';
31
+
32
+/*
33
+ * TODO(jmbrown):
34
+ * In this stack:
35
+ * - Insert useFire import
36
+ * - Assert no lingering fire calls
37
+ * - Ensure a fired function is not called regularly elsewhere in the same effect
38
+ *
39
+ * Future:
40
+ * - rewrite dep arrays
41
+ * - traverse object methods
42
+ * - method calls
43
+ * - React.useEffect calls
44
+ */
45
+
46
+const CANNOT_COMPILE_FIRE = 'Cannot compile `fire`';
47
+
48
+export function transformFire(fn: HIRFunction): void {
49
+ const context = new Context(fn.env);
50
+ replaceFireFunctions(fn, context);
51
+ context.throwIfErrorsFound();
52
+}
53
+
54
+function replaceFireFunctions(fn: HIRFunction, context: Context): void {
55
+ let hasRewrite = false;
56
+ for (const [, block] of fn.body.blocks) {
57
+ const rewriteInstrs = new Map<InstructionId, Array<Instruction>>();
58
+ const deleteInstrs = new Set<InstructionId>();
59
+ for (const instr of block.instructions) {
60
+ const {value, lvalue} = instr;
61
+ if (
62
+ value.kind === 'CallExpression' &&
63
+ isUseEffectHookType(value.callee.identifier) &&
64
+ value.args.length > 0 &&
65
+ value.args[0].kind === 'Identifier'
66
+ ) {
67
+ const lambda = context.getFunctionExpression(
68
+ value.args[0].identifier.id,
69
+ );
70
+ if (lambda != null) {
71
+ const capturedCallees =
72
+ visitFunctionExpressionAndPropagateFireDependencies(
73
+ lambda,
74
+ context,
75
+ true,
76
+ );
77
+
78
+ // Add useFire calls for all fire calls in found in the lambda
79
+ const newInstrs = [];
80
+ for (const [
81
+ fireCalleePlace,
82
+ fireCalleeInfo,
83
+ ] of capturedCallees.entries()) {
84
+ if (!context.hasCalleeWithInsertedFire(fireCalleePlace)) {
85
+ context.addCalleeWithInsertedFire(fireCalleePlace);
86
+ const loadUseFireInstr = makeLoadUseFireInstruction(fn.env);
87
+ const loadFireCalleeInstr = makeLoadFireCalleeInstruction(
88
+ fn.env,
89
+ fireCalleeInfo.capturedCalleeIdentifier,
90
+ );
91
+ const callUseFireInstr = makeCallUseFireInstruction(
92
+ fn.env,
93
+ loadUseFireInstr.lvalue,
94
+ loadFireCalleeInstr.lvalue,
95
+ );
96
+ const storeUseFireInstr = makeStoreUseFireInstruction(
97
+ fn.env,
98
+ callUseFireInstr.lvalue,
99
+ fireCalleeInfo.fireFunctionBinding,
100
+ );
101
+ newInstrs.push(
102
+ loadUseFireInstr,
103
+ loadFireCalleeInstr,
104
+ callUseFireInstr,
105
+ storeUseFireInstr,
106
+ );
107
+
108
+ // We insert all of these instructions before the useEffect is loaded
109
+ const loadUseEffectInstrId = context.getLoadGlobalInstrId(
110
+ value.callee.identifier.id,
111
+ );
112
+ if (loadUseEffectInstrId == null) {
113
+ context.pushError({
114
+ loc: value.loc,
115
+ description: null,
116
+ severity: ErrorSeverity.Invariant,
117
+ reason: '[InsertFire] No LoadGlobal found for useEffect call',
118
+ suggestions: null,
119
+ });
120
+ continue;
121
+ }
122
+ rewriteInstrs.set(loadUseEffectInstrId, newInstrs);
123
+ }
124
+ }
125
+ }
126
+ } else if (
127
+ value.kind === 'CallExpression' &&
128
+ value.callee.identifier.type.kind === 'Function' &&
129
+ value.callee.identifier.type.shapeId === BuiltInFireId &&
130
+ context.inUseEffectLambda()
131
+ ) {
132
+ /*
133
+ * We found a fire(callExpr()) call. We remove the `fire()` call and replace the callExpr()
134
+ * with a freshly generated fire function binding. We'll insert the useFire call before the
135
+ * useEffect call, which happens in the CallExpression (useEffect) case above.
136
+ */
137
+
138
+ /*
139
+ * We only allow fire to be called with a CallExpression: `fire(f())`
140
+ * TODO: add support for method calls: `fire(this.method())`
141
+ */
142
+ if (value.args.length === 1 && value.args[0].kind === 'Identifier') {
143
+ const callExpr = context.getCallExpression(
144
+ value.args[0].identifier.id,
145
+ );
146
+
147
+ if (callExpr != null) {
148
+ const calleeId = callExpr.callee.identifier.id;
149
+ const loadLocal = context.getLoadLocalInstr(calleeId);
150
+ if (loadLocal == null) {
151
+ context.pushError({
152
+ loc: value.loc,
153
+ description: null,
154
+ severity: ErrorSeverity.Invariant,
155
+ reason:
156
+ '[InsertFire] No loadLocal found for fire call argument',
157
+ suggestions: null,
158
+ });
159
+ continue;
160
+ }
161
+
162
+ const fireFunctionBinding =
163
+ context.getOrGenerateFireFunctionBinding(loadLocal.place);
164
+
165
+ loadLocal.place = {...fireFunctionBinding};
166
+
167
+ // Delete the fire call expression
168
+ deleteInstrs.add(instr.id);
169
+ } else {
170
+ context.pushError({
171
+ loc: value.loc,
172
+ description:
173
+ '`fire()` can only receive a function call such as `fire(fn(a,b)). Method calls and other expressions are not allowed',
174
+ severity: ErrorSeverity.InvalidReact,
175
+ reason: CANNOT_COMPILE_FIRE,
176
+ suggestions: null,
177
+ });
178
+ }
179
+ } else {
180
+ let description: string =
181
+ 'fire() can only take in a single call expression as an argument';
182
+ if (value.args.length === 0) {
183
+ description += ' but received none';
184
+ } else if (value.args.length > 1) {
185
+ description += ' but received multiple arguments';
186
+ } else if (value.args[0].kind === 'Spread') {
187
+ description += ' but received a spread argument';
188
+ }
189
+ context.pushError({
190
+ loc: value.loc,
191
+ description,
192
+ severity: ErrorSeverity.InvalidReact,
193
+ reason: CANNOT_COMPILE_FIRE,
194
+ suggestions: null,
195
+ });
196
+ }
197
+ } else if (value.kind === 'CallExpression') {
198
+ context.addCallExpression(lvalue.identifier.id, value);
199
+ } else if (
200
+ value.kind === 'FunctionExpression' &&
201
+ context.inUseEffectLambda()
202
+ ) {
203
+ visitFunctionExpressionAndPropagateFireDependencies(
204
+ value,
205
+ context,
206
+ false,
207
+ );
208
+ } else if (value.kind === 'FunctionExpression') {
209
+ context.addFunctionExpression(lvalue.identifier.id, value);
210
+ } else if (value.kind === 'LoadLocal') {
211
+ context.addLoadLocalInstr(lvalue.identifier.id, value);
212
+ } else if (
213
+ value.kind === 'LoadGlobal' &&
214
+ value.binding.kind === 'ImportSpecifier' &&
215
+ value.binding.module === 'react' &&
216
+ value.binding.imported === 'fire' &&
217
+ context.inUseEffectLambda()
218
+ ) {
219
+ deleteInstrs.add(instr.id);
220
+ } else if (value.kind === 'LoadGlobal') {
221
+ context.addLoadGlobalInstrId(lvalue.identifier.id, instr.id);
222
+ }
223
+ }
224
+ block.instructions = rewriteInstructions(rewriteInstrs, block.instructions);
225
+ block.instructions = deleteInstructions(deleteInstrs, block.instructions);
226
+
227
+ if (rewriteInstrs.size > 0 || deleteInstrs.size > 0) {
228
+ hasRewrite = true;
229
+ }
230
+ }
231
+
232
+ if (hasRewrite) {
233
+ markInstructionIds(fn.body);
234
+ }
235
+}
236
+
237
+/**
238
+ * Traverses a function expression to find fire calls fire(foo()) and replaces them with
239
+ * fireFoo().
240
+ *
241
+ * When a function captures a fire call we need to update its context to reflect the newly created
242
+ * fire function bindings and update the LoadLocals referenced by the function's dependencies.
243
+ *
244
+ * @param isUseEffect is necessary so we can keep track of when we should additionally insert
245
+ * useFire hooks calls.
246
+ */
247
+function visitFunctionExpressionAndPropagateFireDependencies(
248
+ fnExpr: FunctionExpression,
249
+ context: Context,
250
+ enteringUseEffect: boolean,
251
+): FireCalleesToFireFunctionBinding {
252
+ let withScope = enteringUseEffect
253
+ ? context.withUseEffectLambdaScope.bind(context)
254
+ : context.withFunctionScope.bind(context);
255
+
256
+ const calleesCapturedByFnExpression = withScope(() =>
257
+ replaceFireFunctions(fnExpr.loweredFunc.func, context),
258
+ );
259
+
260
+ /*
261
+ * Make a mapping from each dependency to the corresponding LoadLocal for it so that
262
+ * we can replace the loaded place with the generated fire function binding
263
+ */
264
+ const loadLocalsToDepLoads = new Map<IdentifierId, LoadLocal>();
265
+ for (const dep of fnExpr.loweredFunc.dependencies) {
266
+ const loadLocal = context.getLoadLocalInstr(dep.identifier.id);
267
+ if (loadLocal != null) {
268
+ loadLocalsToDepLoads.set(loadLocal.place.identifier.id, loadLocal);
269
+ }
270
+ }
271
+
272
+ const replacedCallees = new Map<IdentifierId, Place>();
273
+ for (const [
274
+ calleeIdentifierId,
275
+ loadedFireFunctionBindingPlace,
276
+ ] of calleesCapturedByFnExpression.entries()) {
277
+ /*
278
+ * Given the ids of captured fire callees, look at the deps for loads of those identifiers
279
+ * and replace them with the new fire function binding
280
+ */
281
+ const loadLocal = loadLocalsToDepLoads.get(calleeIdentifierId);
282
+ if (loadLocal == null) {
283
+ context.pushError({
284
+ loc: fnExpr.loc,
285
+ description: null,
286
+ severity: ErrorSeverity.Invariant,
287
+ reason:
288
+ '[InsertFire] No loadLocal found for fire call argument for lambda',
289
+ suggestions: null,
290
+ });
291
+ continue;
292
+ }
293
+
294
+ const oldPlaceId = loadLocal.place.identifier.id;
295
+ loadLocal.place = {
296
+ ...loadedFireFunctionBindingPlace.fireFunctionBinding,
297
+ };
298
+
299
+ replacedCallees.set(
300
+ oldPlaceId,
301
+ loadedFireFunctionBindingPlace.fireFunctionBinding,
302
+ );
303
+ }
304
+
305
+ // For each replaced callee, update the context of the function expression to track it
306
+ for (
307
+ let contextIdx = 0;
308
+ contextIdx < fnExpr.loweredFunc.func.context.length;
309
+ contextIdx++
310
+ ) {
311
+ const contextItem = fnExpr.loweredFunc.func.context[contextIdx];
312
+ const replacedCallee = replacedCallees.get(contextItem.identifier.id);
313
+ if (replacedCallee != null) {
314
+ fnExpr.loweredFunc.func.context[contextIdx] = replacedCallee;
315
+ }
316
+ }
317
+
318
+ context.mergeCalleesFromInnerScope(calleesCapturedByFnExpression);
319
+
320
+ return calleesCapturedByFnExpression;
321
+}
322
+
323
+function makeLoadUseFireInstruction(env: Environment): Instruction {
324
+ const useFirePlace = createTemporaryPlace(env, GeneratedSource);
325
+ useFirePlace.effect = Effect.Read;
326
+ useFirePlace.identifier.type = DefaultNonmutatingHook;
327
+ const instrValue: InstructionValue = {
328
+ kind: 'LoadGlobal',
329
+ binding: {
330
+ kind: 'ImportSpecifier',
331
+ name: 'useFire',
332
+ module: 'react',
333
+ imported: 'useFire',
334
+ },
335
+ loc: GeneratedSource,
336
+ };
337
+ return {
338
+ id: makeInstructionId(0),
339
+ value: instrValue,
340
+ lvalue: {...useFirePlace},
341
+ loc: GeneratedSource,
342
+ };
343
+}
344
+
345
+function makeLoadFireCalleeInstruction(
346
+ env: Environment,
347
+ fireCalleeIdentifier: Identifier,
348
+): Instruction {
349
+ const loadedFireCallee = createTemporaryPlace(env, GeneratedSource);
350
+ const fireCallee: Place = {
351
+ kind: 'Identifier',
352
+ identifier: fireCalleeIdentifier,
353
+ reactive: false,
354
+ effect: Effect.Unknown,
355
+ loc: fireCalleeIdentifier.loc,
356
+ };
357
+ return {
358
+ id: makeInstructionId(0),
359
+ value: {
360
+ kind: 'LoadLocal',
361
+ loc: GeneratedSource,
362
+ place: {...fireCallee},
363
+ },
364
+ lvalue: {...loadedFireCallee},
365
+ loc: GeneratedSource,
366
+ };
367
+}
368
+
369
+function makeCallUseFireInstruction(
370
+ env: Environment,
371
+ useFirePlace: Place,
372
+ argPlace: Place,
373
+): Instruction {
374
+ const useFireCallResultPlace = createTemporaryPlace(env, GeneratedSource);
375
+ useFireCallResultPlace.effect = Effect.Read;
376
+
377
+ const useFireCall: CallExpression = {
378
+ kind: 'CallExpression',
379
+ callee: {...useFirePlace},
380
+ args: [argPlace],
381
+ loc: GeneratedSource,
382
+ };
383
+
384
+ return {
385
+ id: makeInstructionId(0),
386
+ value: useFireCall,
387
+ lvalue: {...useFireCallResultPlace},
388
+ loc: GeneratedSource,
389
+ };
390
+}
391
+
392
+function makeStoreUseFireInstruction(
393
+ env: Environment,
394
+ useFireCallResultPlace: Place,
395
+ fireFunctionBindingPlace: Place,
396
+): Instruction {
397
+ promoteTemporary(fireFunctionBindingPlace.identifier);
398
+
399
+ const fireFunctionBindingLValuePlace = createTemporaryPlace(
400
+ env,
401
+ GeneratedSource,
402
+ );
403
+ return {
404
+ id: makeInstructionId(0),
405
+ value: {
406
+ kind: 'StoreLocal',
407
+ lvalue: {
408
+ kind: InstructionKind.Const,
409
+ place: {...fireFunctionBindingPlace},
410
+ },
411
+ value: {...useFireCallResultPlace},
412
+ type: null,
413
+ loc: GeneratedSource,
414
+ },
415
+ lvalue: fireFunctionBindingLValuePlace,
416
+ loc: GeneratedSource,
417
+ };
418
+}
419
+
420
+type FireCalleesToFireFunctionBinding = Map<
421
+ IdentifierId,
422
+ {
423
+ fireFunctionBinding: Place;
424
+ capturedCalleeIdentifier: Identifier;
425
+ }
426
+>;
427
+
428
+class Context {
429
+ #env: Environment;
430
+
431
+ #errors: CompilerError = new CompilerError();
432
+
433
+ /*
434
+ * Used to look up the call expression passed to a `fire(callExpr())`. Gives back
435
+ * the `callExpr()`.
436
+ */
437
+ #callExpressions = new Map<IdentifierId, CallExpression>();
438
+
439
+ /*
440
+ * We keep track of function expressions so that we can traverse them when
441
+ * we encounter a lambda passed to a useEffect call
442
+ */
443
+ #functionExpressions = new Map<IdentifierId, FunctionExpression>();
444
+
445
+ /*
446
+ * Mapping from lvalue ids to the LoadLocal for it. Allows us to replace dependency LoadLocals.
447
+ */
448
+ #loadLocals = new Map<IdentifierId, LoadLocal>();
449
+
450
+ /*
451
+ * Maps all of the fire callees found in a component/hook to the generated fire function places
452
+ * we create for them. Allows us to reuse already-inserted useFire results
453
+ */
454
+ #fireCalleesToFireFunctions: Map<IdentifierId, Place> = new Map();
455
+
456
+ /*
457
+ * The callees for which we have already created fire bindings. Used to skip inserting a new
458
+ * useFire call for a fire callee if one has already been created.
459
+ */
460
+ #calleesWithInsertedFire = new Set<IdentifierId>();
461
+
462
+ /*
463
+ * A mapping from fire callees to the created fire function bindings that are reachable from this
464
+ * scope.
465
+ *
466
+ * We additionally keep track of the captured callee identifier so that we can properly reference
467
+ * it in the place where we LoadLocal the callee as an argument to useFire.
468
+ */
469
+ #capturedCalleeIdentifierIds: FireCalleesToFireFunctionBinding = new Map();
470
+
471
+ /*
472
+ * We only transform fire calls if we're syntactically within a useEffect lambda (for now)
473
+ */
474
+ #inUseEffectLambda = false;
475
+
476
+ /*
477
+ * Mapping from useEffect callee identifier ids to the instruction id of the
478
+ * load global instruction for the useEffect call. We use this to insert the
479
+ * useFire calls before the useEffect call
480
+ */
481
+ #loadGlobalInstructionIds = new Map<IdentifierId, InstructionId>();
482
+
483
+ constructor(env: Environment) {
484
+ this.#env = env;
485
+ }
486
+
487
+ pushError(error: CompilerErrorDetailOptions): void {
488
+ this.#errors.push(error);
489
+ }
490
+
491
+ withFunctionScope(fn: () => void): FireCalleesToFireFunctionBinding {
492
+ fn();
493
+ return this.#capturedCalleeIdentifierIds;
494
+ }
495
+
496
+ withUseEffectLambdaScope(fn: () => void): FireCalleesToFireFunctionBinding {
497
+ const capturedCalleeIdentifierIds = this.#capturedCalleeIdentifierIds;
498
+ const inUseEffectLambda = this.#inUseEffectLambda;
499
+
500
+ this.#capturedCalleeIdentifierIds = new Map();
501
+ this.#inUseEffectLambda = true;
502
+
503
+ const resultCapturedCalleeIdentifierIds = this.withFunctionScope(fn);
504
+
505
+ this.#capturedCalleeIdentifierIds = capturedCalleeIdentifierIds;
506
+ this.#inUseEffectLambda = inUseEffectLambda;
507
+
508
+ return resultCapturedCalleeIdentifierIds;
509
+ }
510
+
511
+ addCallExpression(id: IdentifierId, callExpr: CallExpression): void {
512
+ this.#callExpressions.set(id, callExpr);
513
+ }
514
+
515
+ getCallExpression(id: IdentifierId): CallExpression | undefined {
516
+ return this.#callExpressions.get(id);
517
+ }
518
+
519
+ addLoadLocalInstr(id: IdentifierId, loadLocal: LoadLocal): void {
520
+ this.#loadLocals.set(id, loadLocal);
521
+ }
522
+
523
+ getLoadLocalInstr(id: IdentifierId): LoadLocal | undefined {
524
+ return this.#loadLocals.get(id);
525
+ }
526
+
527
+ getOrGenerateFireFunctionBinding(callee: Place): Place {
528
+ const fireFunctionBinding = getOrInsertWith(
529
+ this.#fireCalleesToFireFunctions,
530
+ callee.identifier.id,
531
+ () => createTemporaryPlace(this.#env, GeneratedSource),
532
+ );
533
+
534
+ this.#capturedCalleeIdentifierIds.set(callee.identifier.id, {
535
+ fireFunctionBinding,
536
+ capturedCalleeIdentifier: callee.identifier,
537
+ });
538
+
539
+ return fireFunctionBinding;
540
+ }
541
+
542
+ mergeCalleesFromInnerScope(
543
+ innerCallees: FireCalleesToFireFunctionBinding,
544
+ ): void {
545
+ for (const [id, calleeInfo] of innerCallees.entries()) {
546
+ this.#capturedCalleeIdentifierIds.set(id, calleeInfo);
547
+ }
548
+ }
549
+
550
+ addCalleeWithInsertedFire(id: IdentifierId): void {
551
+ this.#calleesWithInsertedFire.add(id);
552
+ }
553
+
554
+ hasCalleeWithInsertedFire(id: IdentifierId): boolean {
555
+ return this.#calleesWithInsertedFire.has(id);
556
+ }
557
+
558
+ inUseEffectLambda(): boolean {
559
+ return this.#inUseEffectLambda;
560
+ }
561
+
562
+ addFunctionExpression(id: IdentifierId, fn: FunctionExpression): void {
563
+ this.#functionExpressions.set(id, fn);
564
+ }
565
+
566
+ getFunctionExpression(id: IdentifierId): FunctionExpression | undefined {
567
+ return this.#functionExpressions.get(id);
568
+ }
569
+
570
+ addLoadGlobalInstrId(id: IdentifierId, instrId: InstructionId): void {
571
+ this.#loadGlobalInstructionIds.set(id, instrId);
572
+ }
573
+
574
+ getLoadGlobalInstrId(id: IdentifierId): InstructionId | undefined {
575
+ return this.#loadGlobalInstructionIds.get(id);
576
+ }
577
+
578
+ throwIfErrorsFound(): void {
579
+ if (this.#errors.hasErrors()) throw this.#errors;
580
+ }
581
+}
582
+
583
+function deleteInstructions(
584
+ deleteInstrs: Set<InstructionId>,
585
+ instructions: Array<Instruction>,
586
+): Array<Instruction> {
587
+ if (deleteInstrs.size > 0) {
588
+ const newInstrs = instructions.filter(instr => !deleteInstrs.has(instr.id));
589
+ return newInstrs;
590
+ }
591
+ return instructions;
592
+}
593
+
594
+function rewriteInstructions(
595
+ rewriteInstrs: Map<InstructionId, Array<Instruction>>,
596
+ instructions: Array<Instruction>,
597
+): Array<Instruction> {
598
+ if (rewriteInstrs.size > 0) {
599
+ const newInstrs = [];
600
+ for (const instr of instructions) {
601
+ const newInstrsAtId = rewriteInstrs.get(instr.id);
602
+ if (newInstrsAtId != null) {
603
+ newInstrs.push(...newInstrsAtId, instr);
604
+ } else {
605
+ newInstrs.push(instr);
606
+ }
607
+ }
608
+
609
+ return newInstrs;
610
+ }
611
+
612
+ return instructions;
613
+}