diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts index adee97238b2b7..ed11763eee881 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts @@ -265,6 +265,10 @@ export function compileProgram( ); const lintError = suppressionsToCompilerError(suppressions); let hasCriticalError = lintError != null; + const queue: Array<{ + fn: BabelFn; + fnType: ReactFunctionType; + }> = []; const compiledFns: Array = []; const traverseFunction = (fn: BabelFn, pass: CompilerPass): void => { @@ -281,6 +285,47 @@ export function compileProgram( ALREADY_COMPILED.add(fn.node); fn.skip(); + queue.push({ fn, fnType }); + }; + + // Main traversal to compile with Forget + program.traverse( + { + ClassDeclaration(node: NodePath) { + /* + * Don't visit functions defined inside classes, because they + * can reference `this` which is unsafe for compilation + */ + node.skip(); + return; + }, + + ClassExpression(node: NodePath) { + /* + * Don't visit functions defined inside classes, because they + * can reference `this` which is unsafe for compilation + */ + node.skip(); + return; + }, + + FunctionDeclaration: traverseFunction, + + FunctionExpression: traverseFunction, + + ArrowFunctionExpression: traverseFunction, + }, + { + ...pass, + opts: { ...pass.opts, ...pass.opts }, + filename: pass.filename ?? null, + } + ); + + const processFn = ( + fn: BabelFn, + fnType: ReactFunctionType + ): null | CodegenFunction => { if (lintError != null) { /** * Note that Babel does not attach comment nodes to nodes; they are dangling off of the @@ -335,52 +380,33 @@ export function compileProgram( } catch (err) { hasCriticalError ||= isCriticalError(err); handleError(err, pass, fn.node.loc ?? null); - return; + return null; } if (!pass.opts.noEmit && !hasCriticalError) { - compiledFns.push({ originalFn: fn, compiledFn }); + return compiledFn; } + return null; }; - // Main traversal to compile with Forget - program.traverse( - { - ClassDeclaration(node: NodePath) { - /* - * Don't visit functions defined inside classes, because they - * can reference `this` which is unsafe for compilation - */ - node.skip(); - return; - }, - - ClassExpression(node: NodePath) { - /* - * Don't visit functions defined inside classes, because they - * can reference `this` which is unsafe for compilation - */ - node.skip(); - return; - }, - - FunctionDeclaration: traverseFunction, - - FunctionExpression: traverseFunction, - - ArrowFunctionExpression: traverseFunction, - }, - { - ...pass, - opts: { ...pass.opts, ...pass.opts }, - filename: pass.filename ?? null, + while (queue.length !== 0) { + const current = queue.shift()!; + const compiled = processFn(current.fn, current.fnType); + if (compiled === null) { + continue; } - ); + compiledFns.push({ + compiledFn: compiled, + originalFn: current.fn, + }); + } if (pass.opts.gating != null) { const error = checkFunctionReferencedBeforeDeclarationAtTopLevel( program, - compiledFns.map(({ originalFn }) => originalFn) + compiledFns.map((result) => { + return result.originalFn; + }) ); if (error) { handleError(error, pass, null); @@ -439,7 +465,8 @@ export function compileProgram( * Only insert Forget-ified functions if we have not encountered a critical * error elsewhere in the file, regardless of bailout mode. */ - for (const { originalFn, compiledFn } of compiledFns) { + for (const result of compiledFns) { + const { originalFn, compiledFn } = result; const transformedFn = createNewFunctionNode(originalFn, compiledFn); if (gating != null) {