[compiler] Refactor Program to use queue of functions to compile
Refactors Program.ts to first traverse the `Program` node and build up a queue of functions to visit, then iterate that queue and compile the functions. This doesn't change behavior, but allows the next diff to add additional items to the queue during compilation (for function outlining). ghstack-source-id: 858527c30ccc26b3aa6fe75a4746fce0820b316f Pull Request resolved: https://github.com/facebook/react/pull/30330
Joe Savona committed
Jul 15, 2024 at 12:28 UTC
a8c5e23e5ec13a32196bf2512c5cbd2520f7be28
1 file changed
+63
-36
compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts
+63
-36
@@ -265,6 +265,10 @@ export function compileProgram(
265
);
266
const lintError = suppressionsToCompilerError(suppressions);
267
let hasCriticalError = lintError != null;
268
+ const queue: Array<{
269
+ fn: BabelFn;
270
+ fnType: ReactFunctionType;
271
+ }> = [];
272
const compiledFns: Array<CompileResult> = [];
273
274
const traverseFunction = (fn: BabelFn, pass: CompilerPass): void => {
@@ -281,6 +285,47 @@ export function compileProgram(
285
ALREADY_COMPILED.add(fn.node);
286
fn.skip();
287
288
+ queue.push({ fn, fnType });
289
+ };
290
+
291
+ // Main traversal to compile with Forget
292
+ program.traverse(
293
+ {
294
+ ClassDeclaration(node: NodePath<t.ClassDeclaration>) {
295
+ /*
296
+ * Don't visit functions defined inside classes, because they
297
+ * can reference `this` which is unsafe for compilation
298
+ */
299
+ node.skip();
300
+ return;
301
+ },
302
+
303
+ ClassExpression(node: NodePath<t.ClassExpression>) {
304
+ /*
305
+ * Don't visit functions defined inside classes, because they
306
+ * can reference `this` which is unsafe for compilation
307
+ */
308
+ node.skip();
309
+ return;
310
+ },
311
+
312
+ FunctionDeclaration: traverseFunction,
313
+
314
+ FunctionExpression: traverseFunction,
315
+
316
+ ArrowFunctionExpression: traverseFunction,
317
+ },
318
+ {
319
+ ...pass,
320
+ opts: { ...pass.opts, ...pass.opts },
321
+ filename: pass.filename ?? null,
322
+ }
323
+ );
324
+
325
+ const processFn = (
326
+ fn: BabelFn,
327
+ fnType: ReactFunctionType
328
+ ): null | CodegenFunction => {
329
if (lintError != null) {
330
/**
331
* Note that Babel does not attach comment nodes to nodes; they are dangling off of the
@@ -335,52 +380,33 @@ export function compileProgram(
380
} catch (err) {
381
hasCriticalError ||= isCriticalError(err);
382
handleError(err, pass, fn.node.loc ?? null);
338
- return;
383
+ return null;
384
}
385
386
if (!pass.opts.noEmit && !hasCriticalError) {
342
- compiledFns.push({ originalFn: fn, compiledFn });
387
+ return compiledFn;
388
}
389
+ return null;
390
};
391
346
- // Main traversal to compile with Forget
347
- program.traverse(
348
- {
349
- ClassDeclaration(node: NodePath<t.ClassDeclaration>) {
350
- /*
351
- * Don't visit functions defined inside classes, because they
352
- * can reference `this` which is unsafe for compilation
353
- */
354
- node.skip();
355
- return;
356
- },
357
-
358
- ClassExpression(node: NodePath<t.ClassExpression>) {
359
- /*
360
- * Don't visit functions defined inside classes, because they
361
- * can reference `this` which is unsafe for compilation
362
- */
363
- node.skip();
364
- return;
365
- },
366
-
367
- FunctionDeclaration: traverseFunction,
368
-
369
- FunctionExpression: traverseFunction,
370
-
371
- ArrowFunctionExpression: traverseFunction,
372
- },
373
- {
374
- ...pass,
375
- opts: { ...pass.opts, ...pass.opts },
376
- filename: pass.filename ?? null,
392
+ while (queue.length !== 0) {
393
+ const current = queue.shift()!;
394
+ const compiled = processFn(current.fn, current.fnType);
395
+ if (compiled === null) {
396
+ continue;
397
}
378
- );
398
+ compiledFns.push({
399
+ compiledFn: compiled,
400
+ originalFn: current.fn,
401
+ });
402
+ }
403
404
if (pass.opts.gating != null) {
405
const error = checkFunctionReferencedBeforeDeclarationAtTopLevel(
406
program,
383
- compiledFns.map(({ originalFn }) => originalFn)
407
+ compiledFns.map((result) => {
408
+ return result.originalFn;
409
+ })
410
);
411
if (error) {
412
handleError(error, pass, null);
@@ -439,7 +465,8 @@ export function compileProgram(
465
* Only insert Forget-ified functions if we have not encountered a critical
466
* error elsewhere in the file, regardless of bailout mode.
467
*/
442
- for (const { originalFn, compiledFn } of compiledFns) {
468
+ for (const result of compiledFns) {
469
+ const { originalFn, compiledFn } = result;
470
const transformedFn = createNewFunctionNode(originalFn, compiledFn);
471
472
if (gating != null) {