@samitouri / QOS-React-2 / commits / 5069e18060

[compiler][be] Make program traversal more readable (#33147)

React Compiler's program traversal logic is pretty lengthy and complex as we've added a lot of features piecemeal. `compileProgram` is 300+ lines long and has confusing control flow (defining helpers inline, invoking visitors, mutating-asts-while-iterating, mutating global `ALREADY_COMPILED` state). - Moved more stuff to `ProgramContext` - Separated `compileProgram` into a bunch of helpers Tested by syncing this stack to a Meta codebase and observing no compilation output changes (D74487851, P1806855669, P1806855379) --- [//]: # (BEGIN SAPLING FOOTER) Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/facebook/react/pull/33147). * #33149 * #33148 * __->__ #33147

mofeiZ committed May 9, 2025 at 13:23 UTC 5069e18060e00d7c07b2b04ebc8a3fa21e2d810a
9 files changed +521 -250
compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Imports.ts
+49 -13
@@ -18,8 +18,9 @@ import {
18 import {getOrInsertWith} from '../Utils/utils';
19 import {ExternalFunction, isHookName} from '../HIR/Environment';
20 import {Err, Ok, Result} from '../Utils/Result';
21 -import {CompilerReactTarget} from './Options';
22 -import {getReactCompilerRuntimeModule} from './Program';
21 +import {LoggerEvent, PluginOptions} from './Options';
22 +import {BabelFn, getReactCompilerRuntimeModule} from './Program';
23 +import {SuppressionRange} from './Suppression';
24
25 export function validateRestrictedImports(
26 path: NodePath<t.Program>,
@@ -52,32 +53,61 @@ export function validateRestrictedImports(
53 }
54 }
55
56 +type ProgramContextOptions = {
57 + program: NodePath<t.Program>;
58 + suppressions: Array<SuppressionRange>;
59 + opts: PluginOptions;
60 + filename: string | null;
61 + code: string | null;
62 +};
63 export class ProgramContext {
56 - /* Program and environment context */
64 + /**
65 + * Program and environment context
66 + */
67 scope: BabelScope;
68 + opts: PluginOptions;
69 + filename: string | null;
70 + code: string | null;
71 reactRuntimeModule: string;
59 - hookPattern: string | null;
72 + suppressions: Array<SuppressionRange>;
73
74 + /*
75 + * This is a hack to work around what seems to be a Babel bug. Babel doesn't
76 + * consistently respect the `skip()` function to avoid revisiting a node within
77 + * a pass, so we use this set to track nodes that we have compiled.
78 + */
79 + alreadyCompiled: WeakSet<object> | Set<object> = new (WeakSet ?? Set)();
80 // known generated or referenced identifiers in the program
81 knownReferencedNames: Set<string> = new Set();
82 // generated imports
83 imports: Map<string, Map<string, NonLocalImportSpecifier>> = new Map();
84
66 - constructor(
67 - program: NodePath<t.Program>,
68 - reactRuntimeModule: CompilerReactTarget,
69 - hookPattern: string | null,
70 - ) {
71 - this.hookPattern = hookPattern;
85 + /**
86 + * Metadata from compilation
87 + */
88 + retryErrors: Array<{fn: BabelFn; error: CompilerError}> = [];
89 + inferredEffectLocations: Set<t.SourceLocation> = new Set();
90 +
91 + constructor({
92 + program,
93 + suppressions,
94 + opts,
95 + filename,
96 + code,
97 + }: ProgramContextOptions) {
98 this.scope = program.scope;
73 - this.reactRuntimeModule = getReactCompilerRuntimeModule(reactRuntimeModule);
99 + this.opts = opts;
100 + this.filename = filename;
101 + this.code = code;
102 + this.reactRuntimeModule = getReactCompilerRuntimeModule(opts.target);
103 + this.suppressions = suppressions;
104 }
105
106 isHookName(name: string): boolean {
77 - if (this.hookPattern == null) {
107 + if (this.opts.environment.hookPattern == null) {
108 return isHookName(name);
109 } else {
80 - const match = new RegExp(this.hookPattern).exec(name);
110 + const match = new RegExp(this.opts.environment.hookPattern).exec(name);
111 return (
112 match != null && typeof match[1] === 'string' && isHookName(match[1])
113 );
@@ -179,6 +209,12 @@ export class ProgramContext {
209 });
210 return Err(error);
211 }
212 +
213 + logEvent(event: LoggerEvent): void {
214 + if (this.opts.logger != null) {
215 + this.opts.logger.logEvent(this.filename, event);
216 + }
217 + }
218 }
219
220 function getExistingImports(
compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts
+305 -234
@@ -12,7 +12,7 @@ import {
12 CompilerErrorDetail,
13 ErrorSeverity,
14 } from '../CompilerError';
15 -import {EnvironmentConfig, ReactFunctionType} from '../HIR/Environment';
15 +import {ReactFunctionType} from '../HIR/Environment';
16 import {CodegenFunction} from '../ReactiveScopes';
17 import {isComponentDeclaration} from '../Utils/ComponentDeclaration';
18 import {isHookDeclaration} from '../Utils/HookDeclaration';
@@ -43,17 +43,21 @@ export const OPT_OUT_DIRECTIVES = new Set(['use no forget', 'use no memo']);
43
44 export function findDirectiveEnablingMemoization(
45 directives: Array<t.Directive>,
46 -): Array<t.Directive> {
47 - return directives.filter(directive =>
48 - OPT_IN_DIRECTIVES.has(directive.value.value),
46 +): t.Directive | null {
47 + return (
48 + directives.find(directive =>
49 + OPT_IN_DIRECTIVES.has(directive.value.value),
50 + ) ?? null
51 );
52 }
53
54 export function findDirectiveDisablingMemoization(
55 directives: Array<t.Directive>,
54 -): Array<t.Directive> {
55 - return directives.filter(directive =>
56 - OPT_OUT_DIRECTIVES.has(directive.value.value),
56 +): t.Directive | null {
57 + return (
58 + directives.find(directive =>
59 + OPT_OUT_DIRECTIVES.has(directive.value.value),
60 + ) ?? null
61 );
62 }
63
@@ -88,13 +92,16 @@ export type CompileResult = {
92
93 function logError(
94 err: unknown,
91 - pass: CompilerPass,
95 + context: {
96 + opts: PluginOptions;
97 + filename: string | null;
98 + },
99 fnLoc: t.SourceLocation | null,
100 ): void {
94 - if (pass.opts.logger) {
101 + if (context.opts.logger) {
102 if (err instanceof CompilerError) {
103 for (const detail of err.details) {
97 - pass.opts.logger.logEvent(pass.filename, {
104 + context.opts.logger.logEvent(context.filename, {
105 kind: 'CompileError',
106 fnLoc,
107 detail: detail.options,
@@ -108,7 +115,7 @@ function logError(
115 stringifiedError = err?.toString() ?? '[ null ]';
116 }
117
111 - pass.opts.logger.logEvent(pass.filename, {
118 + context.opts.logger.logEvent(context.filename, {
119 kind: 'PipelineError',
120 fnLoc,
121 data: stringifiedError,
@@ -118,13 +125,17 @@ function logError(
125 }
126 function handleError(
127 err: unknown,
121 - pass: CompilerPass,
128 + context: {
129 + opts: PluginOptions;
130 + filename: string | null;
131 + },
132 fnLoc: t.SourceLocation | null,
133 ): void {
124 - logError(err, pass, fnLoc);
134 + logError(err, context, fnLoc);
135 if (
126 - pass.opts.panicThreshold === 'all_errors' ||
127 - (pass.opts.panicThreshold === 'critical_errors' && isCriticalError(err)) ||
136 + context.opts.panicThreshold === 'all_errors' ||
137 + (context.opts.panicThreshold === 'critical_errors' &&
138 + isCriticalError(err)) ||
139 isConfigError(err) // Always throws regardless of panic threshold
140 ) {
141 throw err;
@@ -187,7 +198,6 @@ export function createNewFunctionNode(
198 }
199 }
200 // Avoid visiting the new transformed version
190 - ALREADY_COMPILED.add(transformedFn);
201 return transformedFn;
202 }
203
@@ -239,13 +249,6 @@ function insertNewOutlinedFunctionNode(
249 }
250 }
251
242 -/*
243 - * This is a hack to work around what seems to be a Babel bug. Babel doesn't
244 - * consistently respect the `skip()` function to avoid revisiting a node within
245 - * a pass, so we use this set to track nodes that we have compiled.
246 - */
247 -const ALREADY_COMPILED: WeakSet<object> | Set<object> = new (WeakSet ?? Set)();
248 -
252 const DEFAULT_ESLINT_SUPPRESSIONS = [
253 'react-hooks/exhaustive-deps',
254 'react-hooks/rules-of-hooks',
@@ -268,41 +271,43 @@ function isFilePartOfSources(
271 return false;
272 }
273
271 -export type CompileProgramResult = {
274 +export type CompileProgramMetadata = {
275 retryErrors: Array<{fn: BabelFn; error: CompilerError}>;
276 inferredEffectLocations: Set<t.SourceLocation>;
277 };
278 /**
276 - * `compileProgram` is directly invoked by the react-compiler babel plugin, so
277 - * exceptions thrown by this function will fail the babel build.
278 - * - call `handleError` if your error is recoverable.
279 - * Unless the error is a warning / info diagnostic, compilation of a function
280 - * / entire file should also be skipped.
281 - * - throw an exception if the error is fatal / not recoverable.
282 - * Examples of this are invalid compiler configs or failure to codegen outlined
283 - * functions *after* already emitting optimized components / hooks that invoke
284 - * the outlined functions.
279 + * Main entrypoint for React Compiler.
280 + *
281 + * @param program The Babel program node to compile
282 + * @param pass Compiler configuration and context
283 + * @returns Compilation results or null if compilation was skipped
284 */
285 export function compileProgram(
286 program: NodePath<t.Program>,
287 pass: CompilerPass,
289 -): CompileProgramResult | null {
288 +): CompileProgramMetadata | null {
289 + /**
290 + * This is directly invoked by the react-compiler babel plugin, so exceptions
291 + * thrown by this function will fail the babel build.
292 + * - call `handleError` if your error is recoverable.
293 + * Unless the error is a warning / info diagnostic, compilation of a function
294 + * / entire file should also be skipped.
295 + * - throw an exception if the error is fatal / not recoverable.
296 + * Examples of this are invalid compiler configs or failure to codegen outlined
297 + * functions *after* already emitting optimized components / hooks that invoke
298 + * the outlined functions.
299 + */
300 if (shouldSkipCompilation(program, pass)) {
301 return null;
302 }
293 -
294 - const environment = pass.opts.environment;
295 - const restrictedImportsErr = validateRestrictedImports(program, environment);
303 + const restrictedImportsErr = validateRestrictedImports(
304 + program,
305 + pass.opts.environment,
306 + );
307 if (restrictedImportsErr) {
308 handleError(restrictedImportsErr, pass, null);
309 return null;
310 }
300 -
301 - const programContext = new ProgramContext(
302 - program,
303 - pass.opts.target,
304 - environment.hookPattern,
305 - );
311 /*
312 * Record lint errors and critical errors as depending on Forget's config,
313 * we may still need to run Forget's analysis on every function (even if we
@@ -313,16 +318,88 @@ export function compileProgram(
318 pass.opts.eslintSuppressionRules ?? DEFAULT_ESLINT_SUPPRESSIONS,
319 pass.opts.flowSuppressions,
320 );
316 - const queue: Array<{
317 - kind: 'original' | 'outlined';
318 - fn: BabelFn;
319 - fnType: ReactFunctionType;
320 - }> = [];
321 +
322 + const programContext = new ProgramContext({
323 + program: program,
324 + opts: pass.opts,
325 + filename: pass.filename,
326 + code: pass.code,
327 + suppressions,
328 + });
329 +
330 + const queue: Array<CompileSource> = findFunctionsToCompile(
331 + program,
332 + pass,
333 + programContext,
334 + );
335 const compiledFns: Array<CompileResult> = [];
336
337 + while (queue.length !== 0) {
338 + const current = queue.shift()!;
339 + const compiled = processFn(current.fn, current.fnType, programContext);
340 +
341 + if (compiled != null) {
342 + for (const outlined of compiled.outlined) {
343 + CompilerError.invariant(outlined.fn.outlined.length === 0, {
344 + reason: 'Unexpected nested outlined functions',
345 + loc: outlined.fn.loc,
346 + });
347 + const fn = insertNewOutlinedFunctionNode(
348 + program,
349 + current.fn,
350 + outlined.fn,
351 + );
352 + fn.skip();
353 + programContext.alreadyCompiled.add(fn.node);
354 + if (outlined.type !== null) {
355 + queue.push({
356 + kind: 'outlined',
357 + fn,
358 + fnType: outlined.type,
359 + });
360 + }
361 + }
362 + compiledFns.push({
363 + kind: current.kind,
364 + originalFn: current.fn,
365 + compiledFn: compiled,
366 + });
367 + }
368 + }
369 +
370 + // Avoid modifying the program if we find a program level opt-out
371 + if (findDirectiveDisablingMemoization(program.node.directives) != null) {
372 + return null;
373 + }
374 +
375 + // Insert React Compiler generated functions into the Babel AST
376 + applyCompiledFunctions(program, compiledFns, pass, programContext);
377 +
378 + return {
379 + retryErrors: programContext.retryErrors,
380 + inferredEffectLocations: programContext.inferredEffectLocations,
381 + };
382 +}
383 +
384 +type CompileSource = {
385 + kind: 'original' | 'outlined';
386 + fn: BabelFn;
387 + fnType: ReactFunctionType;
388 +};
389 +/**
390 + * Find all React components and hooks that need to be compiled
391 + *
392 + * @returns An array of React functions from @param program to transform
393 + */
394 +function findFunctionsToCompile(
395 + program: NodePath<t.Program>,
396 + pass: CompilerPass,
397 + programContext: ProgramContext,
398 +): Array<CompileSource> {
399 + const queue: Array<CompileSource> = [];
400 const traverseFunction = (fn: BabelFn, pass: CompilerPass): void => {
324 - const fnType = getReactFunctionType(fn, pass, environment);
325 - if (fnType === null || ALREADY_COMPILED.has(fn.node)) {
401 + const fnType = getReactFunctionType(fn, pass);
402 + if (fnType === null || programContext.alreadyCompiled.has(fn.node)) {
403 return;
404 }
405
@@ -331,7 +408,7 @@ export function compileProgram(
408 * traversal will loop infinitely.
409 * Ensure we avoid visiting the original function again.
410 */
334 - ALREADY_COMPILED.add(fn.node);
411 + programContext.alreadyCompiled.add(fn.node);
412 fn.skip();
413
414 queue.push({kind: 'original', fn, fnType});
@@ -346,7 +423,6 @@ export function compileProgram(
423 * can reference `this` which is unsafe for compilation
424 */
425 node.skip();
349 - return;
426 },
427
428 ClassExpression(node: NodePath<t.ClassExpression>) {
@@ -355,7 +431,6 @@ export function compileProgram(
431 * can reference `this` which is unsafe for compilation
432 */
433 node.skip();
358 - return;
434 },
435
436 FunctionDeclaration: traverseFunction,
@@ -370,205 +445,205 @@ export function compileProgram(
445 filename: pass.filename ?? null,
446 },
447 );
373 - const retryErrors: Array<{fn: BabelFn; error: CompilerError}> = [];
374 - const inferredEffectLocations = new Set<t.SourceLocation>();
375 - const processFn = (
376 - fn: BabelFn,
377 - fnType: ReactFunctionType,
378 - ): null | CodegenFunction => {
379 - let optInDirectives: Array<t.Directive> = [];
380 - let optOutDirectives: Array<t.Directive> = [];
381 - if (fn.node.body.type === 'BlockStatement') {
382 - optInDirectives = findDirectiveEnablingMemoization(
383 - fn.node.body.directives,
384 - );
385 - optOutDirectives = findDirectiveDisablingMemoization(
386 - fn.node.body.directives,
387 - );
388 - }
448 + return queue;
449 +}
450
390 - /**
391 - * Note that Babel does not attach comment nodes to nodes; they are dangling off of the
392 - * Program node itself. We need to figure out whether an eslint suppression range
393 - * applies to this function first.
394 - */
395 - const suppressionsInFunction = filterSuppressionsThatAffectFunction(
396 - suppressions,
397 - fn,
398 - );
399 - let compileResult:
400 - | {kind: 'compile'; compiledFn: CodegenFunction}
401 - | {kind: 'error'; error: unknown};
402 - if (suppressionsInFunction.length > 0) {
403 - compileResult = {
404 - kind: 'error',
405 - error: suppressionsToCompilerError(suppressionsInFunction),
406 - };
407 - } else {
408 - try {
409 - compileResult = {
410 - kind: 'compile',
411 - compiledFn: compileFn(
412 - fn,
413 - environment,
414 - fnType,
415 - 'all_features',
416 - programContext,
417 - pass.opts.logger,
418 - pass.filename,
419 - pass.code,
420 - ),
421 - };
422 - } catch (err) {
423 - compileResult = {kind: 'error', error: err};
424 - }
425 - }
451 +/**
452 + * Try to compile a source function, taking into account all local suppressions,
453 + * opt-ins, and opt-outs.
454 + *
455 + * Errors encountered during compilation are either logged (if recoverable) or
456 + * thrown (if non-recoverable).
457 + *
458 + * @returns the compiled function or null if the function was skipped (due to
459 + * config settings and/or outputs)
460 + */
461 +function processFn(
462 + fn: BabelFn,
463 + fnType: ReactFunctionType,
464 + programContext: ProgramContext,
465 +): null | CodegenFunction {
466 + let directives;
467 + if (fn.node.body.type !== 'BlockStatement') {
468 + directives = {optIn: null, optOut: null};
469 + } else {
470 + directives = {
471 + optIn: findDirectiveEnablingMemoization(fn.node.body.directives),
472 + optOut: findDirectiveDisablingMemoization(fn.node.body.directives),
473 + };
474 + }
475
427 - if (compileResult.kind === 'error') {
428 - /**
429 - * If an opt out directive is present, log only instead of throwing and don't mark as
430 - * containing a critical error.
431 - */
432 - if (optOutDirectives.length > 0) {
433 - logError(compileResult.error, pass, fn.node.loc ?? null);
434 - } else {
435 - handleError(compileResult.error, pass, fn.node.loc ?? null);
436 - }
437 - // If non-memoization features are enabled, retry regardless of error kind
438 - if (
439 - !(environment.enableFire || environment.inferEffectDependencies != null)
440 - ) {
441 - return null;
442 - }
443 - try {
444 - compileResult = {
445 - kind: 'compile',
446 - compiledFn: compileFn(
447 - fn,
448 - environment,
449 - fnType,
450 - 'no_inferred_memo',
451 - programContext,
452 - pass.opts.logger,
453 - pass.filename,
454 - pass.code,
455 - ),
456 - };
457 - if (
458 - !compileResult.compiledFn.hasFireRewrite &&
459 - !compileResult.compiledFn.hasInferredEffect
460 - ) {
461 - return null;
462 - }
463 - } catch (err) {
464 - // TODO: we might want to log error here, but this will also result in duplicate logging
465 - if (err instanceof CompilerError) {
466 - retryErrors.push({fn, error: err});
467 - }
468 - return null;
469 - }
476 + let compiledFn: CodegenFunction;
477 + const compileResult = tryCompileFunction(fn, fnType, programContext);
478 + if (compileResult.kind === 'error') {
479 + if (directives.optOut != null) {
480 + logError(compileResult.error, programContext, fn.node.loc ?? null);
481 + } else {
482 + handleError(compileResult.error, programContext, fn.node.loc ?? null);
483 }
471 -
472 - /**
473 - * Otherwise if 'use no forget/memo' is present, we still run the code through the compiler
474 - * for validation but we don't mutate the babel AST. This allows us to flag if there is an
475 - * unused 'use no forget/memo' directive.
476 - */
477 - if (pass.opts.ignoreUseNoForget === false && optOutDirectives.length > 0) {
478 - for (const directive of optOutDirectives) {
479 - pass.opts.logger?.logEvent(pass.filename, {
480 - kind: 'CompileSkip',
481 - fnLoc: fn.node.body.loc ?? null,
482 - reason: `Skipped due to '${directive.value.value}' directive.`,
483 - loc: directive.loc ?? null,
484 - });
485 - }
484 + const retryResult = retryCompileFunction(fn, fnType, programContext);
485 + if (retryResult == null) {
486 return null;
487 }
488 + compiledFn = retryResult;
489 + } else {
490 + compiledFn = compileResult.compiledFn;
491 + }
492
489 - pass.opts.logger?.logEvent(pass.filename, {
490 - kind: 'CompileSuccess',
491 - fnLoc: fn.node.loc ?? null,
492 - fnName: compileResult.compiledFn.id?.name ?? null,
493 - memoSlots: compileResult.compiledFn.memoSlotsUsed,
494 - memoBlocks: compileResult.compiledFn.memoBlocks,
495 - memoValues: compileResult.compiledFn.memoValues,
496 - prunedMemoBlocks: compileResult.compiledFn.prunedMemoBlocks,
497 - prunedMemoValues: compileResult.compiledFn.prunedMemoValues,
493 + /**
494 + * Otherwise if 'use no forget/memo' is present, we still run the code through the compiler
495 + * for validation but we don't mutate the babel AST. This allows us to flag if there is an
496 + * unused 'use no forget/memo' directive.
497 + */
498 + if (
499 + programContext.opts.ignoreUseNoForget === false &&
500 + directives.optOut != null
501 + ) {
502 + programContext.logEvent({
503 + kind: 'CompileSkip',
504 + fnLoc: fn.node.body.loc ?? null,
505 + reason: `Skipped due to '${directives.optOut.value}' directive.`,
506 + loc: directives.optOut.loc ?? null,
507 });
508 + return null;
509 + }
510 + programContext.logEvent({
511 + kind: 'CompileSuccess',
512 + fnLoc: fn.node.loc ?? null,
513 + fnName: compiledFn.id?.name ?? null,
514 + memoSlots: compiledFn.memoSlotsUsed,
515 + memoBlocks: compiledFn.memoBlocks,
516 + memoValues: compiledFn.memoValues,
517 + prunedMemoBlocks: compiledFn.prunedMemoBlocks,
518 + prunedMemoValues: compiledFn.prunedMemoValues,
519 + });
520
521 + /**
522 + * Always compile functions with opt in directives.
523 + */
524 + if (directives.optIn != null) {
525 + return compiledFn;
526 + } else if (programContext.opts.compilationMode === 'annotation') {
527 /**
501 - * Always compile functions with opt in directives.
528 + * If no opt-in directive is found and the compiler is configured in
529 + * annotation mode, don't insert the compiled function.
530 */
503 - if (optInDirectives.length > 0) {
504 - return compileResult.compiledFn;
505 - } else if (pass.opts.compilationMode === 'annotation') {
506 - /**
507 - * No opt-in directive in annotation mode, so don't insert the compiled function.
508 - */
509 - return null;
510 - }
511 -
512 - if (!pass.opts.noEmit) {
513 - return compileResult.compiledFn;
514 - }
531 + return null;
532 + } else if (programContext.opts.noEmit) {
533 /**
534 * inferEffectDependencies + noEmit is currently only used for linting. In
535 * this mode, add source locations for where the compiler *can* infer effect
536 * dependencies.
537 */
520 - for (const loc of compileResult.compiledFn.inferredEffectLocations) {
521 - if (loc !== GeneratedSource) inferredEffectLocations.add(loc);
522 - }
523 - return null;
524 - };
525 -
526 - while (queue.length !== 0) {
527 - const current = queue.shift()!;
528 - const compiled = processFn(current.fn, current.fnType);
529 - if (compiled === null) {
530 - continue;
531 - }
532 - for (const outlined of compiled.outlined) {
533 - CompilerError.invariant(outlined.fn.outlined.length === 0, {
534 - reason: 'Unexpected nested outlined functions',
535 - loc: outlined.fn.loc,
536 - });
537 - const fn = insertNewOutlinedFunctionNode(
538 - program,
539 - current.fn,
540 - outlined.fn,
541 - );
542 - fn.skip();
543 - ALREADY_COMPILED.add(fn.node);
544 - if (outlined.type !== null) {
545 - queue.push({
546 - kind: 'outlined',
547 - fn,
548 - fnType: outlined.type,
549 - });
538 + for (const loc of compiledFn.inferredEffectLocations) {
539 + if (loc !== GeneratedSource) {
540 + programContext.inferredEffectLocations.add(loc);
541 }
542 }
552 - compiledFns.push({
553 - kind: current.kind,
554 - compiledFn: compiled,
555 - originalFn: current.fn,
556 - });
543 + return null;
544 + } else {
545 + return compiledFn;
546 }
547 +}
548
549 +function tryCompileFunction(
550 + fn: BabelFn,
551 + fnType: ReactFunctionType,
552 + programContext: ProgramContext,
553 +):
554 + | {kind: 'compile'; compiledFn: CodegenFunction}
555 + | {kind: 'error'; error: unknown} {
556 /**
560 - * Do not modify source if there is a module scope level opt out directive.
557 + * Note that Babel does not attach comment nodes to nodes; they are dangling off of the
558 + * Program node itself. We need to figure out whether an eslint suppression range
559 + * applies to this function first.
560 */
562 - const moduleScopeOptOutDirectives = findDirectiveDisablingMemoization(
563 - program.node.directives,
561 + const suppressionsInFunction = filterSuppressionsThatAffectFunction(
562 + programContext.suppressions,
563 + fn,
564 );
565 - if (moduleScopeOptOutDirectives.length > 0) {
565 + if (suppressionsInFunction.length > 0) {
566 + return {
567 + kind: 'error',
568 + error: suppressionsToCompilerError(suppressionsInFunction),
569 + };
570 + }
571 +
572 + try {
573 + return {
574 + kind: 'compile',
575 + compiledFn: compileFn(
576 + fn,
577 + programContext.opts.environment,
578 + fnType,
579 + 'all_features',
580 + programContext,
581 + programContext.opts.logger,
582 + programContext.filename,
583 + programContext.code,
584 + ),
585 + };
586 + } catch (err) {
587 + return {kind: 'error', error: err};
588 + }
589 +}
590 +
591 +/**
592 + * If non-memo feature flags are enabled, retry compilation with a more minimal
593 + * feature set.
594 + *
595 + * @returns a CodegenFunction if retry was successful
596 + */
597 +function retryCompileFunction(
598 + fn: BabelFn,
599 + fnType: ReactFunctionType,
600 + programContext: ProgramContext,
601 +): CodegenFunction | null {
602 + const environment = programContext.opts.environment;
603 + if (
604 + !(environment.enableFire || environment.inferEffectDependencies != null)
605 + ) {
606 return null;
607 }
568 - /*
569 - * Only insert Forget-ified functions if we have not encountered a critical
570 - * error elsewhere in the file, regardless of bailout mode.
608 + /**
609 + * Note that function suppressions are not checked in the retry pipeline, as
610 + * they only affect auto-memoization features.
611 */
612 + try {
613 + const retryResult = compileFn(
614 + fn,
615 + environment,
616 + fnType,
617 + 'no_inferred_memo',
618 + programContext,
619 + programContext.opts.logger,
620 + programContext.filename,
621 + programContext.code,
622 + );
623 +
624 + if (!retryResult.hasFireRewrite && !retryResult.hasInferredEffect) {
625 + return null;
626 + }
627 + return retryResult;
628 + } catch (err) {
629 + // TODO: we might want to log error here, but this will also result in duplicate logging
630 + if (err instanceof CompilerError) {
631 + programContext.retryErrors.push({fn, error: err});
632 + }
633 + return null;
634 + }
635 +}
636 +
637 +/**
638 + * Applies React Compiler generated functions to the babel AST by replacing
639 + * existing functions in place or inserting new declarations.
640 + */
641 +function applyCompiledFunctions(
642 + program: NodePath<t.Program>,
643 + compiledFns: Array<CompileResult>,
644 + pass: CompilerPass,
645 + programContext: ProgramContext,
646 +): void {
647 const referencedBeforeDeclared =
648 pass.opts.gating != null
649 ? getFunctionReferencedBeforeDeclarationAtTopLevel(program, compiledFns)
@@ -576,6 +651,7 @@ export function compileProgram(
651 for (const result of compiledFns) {
652 const {kind, originalFn, compiledFn} = result;
653 const transformedFn = createNewFunctionNode(originalFn, compiledFn);
654 + programContext.alreadyCompiled.add(transformedFn);
655
656 if (referencedBeforeDeclared != null && kind === 'original') {
657 CompilerError.invariant(pass.opts.gating != null, {
@@ -598,7 +674,6 @@ export function compileProgram(
674 if (compiledFns.length > 0) {
675 addImportsToProgram(program, programContext);
676 }
601 - return {retryErrors, inferredEffectLocations};
677 }
678
679 function shouldSkipCompilation(
@@ -640,14 +715,10 @@ function shouldSkipCompilation(
715 function getReactFunctionType(
716 fn: BabelFn,
717 pass: CompilerPass,
643 - /**
644 - * TODO(mofeiZ): remove once we validate PluginOptions with Zod
645 - */
646 - environment: EnvironmentConfig,
718 ): ReactFunctionType | null {
648 - const hookPattern = environment.hookPattern;
719 + const hookPattern = pass.opts.environment.hookPattern;
720 if (fn.node.body.type === 'BlockStatement') {
650 - if (findDirectiveEnablingMemoization(fn.node.body.directives).length > 0)
721 + if (findDirectiveEnablingMemoization(fn.node.body.directives) != null)
722 return getComponentOrHookLike(fn, hookPattern) ?? 'Other';
723 }
724
compiler/packages/babel-plugin-react-compiler/src/Entrypoint/ValidateNoUntransformedReferences.ts
+3 -3
@@ -18,7 +18,7 @@ import {
18 import {getOrInsertWith} from '../Utils/utils';
19 import {Environment} from '../HIR';
20 import {DEFAULT_EXPORT} from '../HIR/Environment';
21 -import {CompileProgramResult} from './Program';
21 +import {CompileProgramMetadata} from './Program';
22
23 function throwInvalidReact(
24 options: Omit<CompilerErrorDetailOptions, 'severity'>,
@@ -109,7 +109,7 @@ export default function validateNoUntransformedReferences(
109 filename: string | null,
110 logger: Logger | null,
111 env: EnvironmentConfig,
112 - compileResult: CompileProgramResult | null,
112 + compileResult: CompileProgramMetadata | null,
113 ): void {
114 const moduleLoadChecks = new Map<
115 string,
@@ -236,7 +236,7 @@ function transformProgram(
236 moduleLoadChecks: Map<string, Map<string, CheckInvalidReferenceFn>>,
237 filename: string | null,
238 logger: Logger | null,
239 - compileResult: CompileProgramResult | null,
239 + compileResult: CompileProgramMetadata | null,
240 ): void {
241 const traversalState: TraversalState = {
242 shouldInvalidateScopes: true,
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/no-emit/no-emit-lint-repro.expect.md renamed
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/no-emit/no-emit-lint-repro.js renamed
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/no-emit/retry-no-emit.expect.md new
+64
@@ -0,0 +1,64 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @inferEffectDependencies @noEmit @panicThreshold:"none" @loggerTestOnly
6 +import {print} from 'shared-runtime';
7 +import useEffectWrapper from 'useEffectWrapper';
8 +
9 +function Foo({propVal}) {
10 + const arr = [propVal];
11 + useEffectWrapper(() => print(arr));
12 +
13 + const arr2 = [];
14 + useEffectWrapper(() => arr2.push(propVal));
15 + arr2.push(2);
16 + return {arr, arr2};
17 +}
18 +
19 +export const FIXTURE_ENTRYPOINT = {
20 + fn: Foo,
21 + params: [{propVal: 1}],
22 + sequentialRenders: [{propVal: 1}, {propVal: 2}],
23 +};
24 +
25 +```
26 +
27 +## Code
28 +
29 +```javascript
30 +// @inferEffectDependencies @noEmit @panicThreshold:"none" @loggerTestOnly
31 +import { print } from "shared-runtime";
32 +import useEffectWrapper from "useEffectWrapper";
33 +
34 +function Foo({ propVal }) {
35 + const arr = [propVal];
36 + useEffectWrapper(() => print(arr));
37 +
38 + const arr2 = [];
39 + useEffectWrapper(() => arr2.push(propVal));
40 + arr2.push(2);
41 + return { arr, arr2 };
42 +}
43 +
44 +export const FIXTURE_ENTRYPOINT = {
45 + fn: Foo,
46 + params: [{ propVal: 1 }],
47 + sequentialRenders: [{ propVal: 1 }, { propVal: 2 }],
48 +};
49 +
50 +```
51 +
52 +## Logs
53 +
54 +```
55 +{"kind":"CompileError","fnLoc":{"start":{"line":5,"column":0,"index":163},"end":{"line":13,"column":1,"index":357},"filename":"retry-no-emit.ts"},"detail":{"reason":"This mutates a variable that React considers immutable","description":null,"loc":{"start":{"line":11,"column":2,"index":320},"end":{"line":11,"column":6,"index":324},"filename":"retry-no-emit.ts","identifierName":"arr2"},"suggestions":null,"severity":"InvalidReact"}}
56 +{"kind":"AutoDepsDecorations","fnLoc":{"start":{"line":7,"column":2,"index":216},"end":{"line":7,"column":36,"index":250},"filename":"retry-no-emit.ts"},"decorations":[{"start":{"line":7,"column":31,"index":245},"end":{"line":7,"column":34,"index":248},"filename":"retry-no-emit.ts","identifierName":"arr"}]}
57 +{"kind":"AutoDepsDecorations","fnLoc":{"start":{"line":10,"column":2,"index":274},"end":{"line":10,"column":44,"index":316},"filename":"retry-no-emit.ts"},"decorations":[{"start":{"line":10,"column":25,"index":297},"end":{"line":10,"column":29,"index":301},"filename":"retry-no-emit.ts","identifierName":"arr2"},{"start":{"line":10,"column":25,"index":297},"end":{"line":10,"column":29,"index":301},"filename":"retry-no-emit.ts","identifierName":"arr2"},{"start":{"line":10,"column":35,"index":307},"end":{"line":10,"column":42,"index":314},"filename":"retry-no-emit.ts","identifierName":"propVal"}]}
58 +{"kind":"CompileSuccess","fnLoc":{"start":{"line":5,"column":0,"index":163},"end":{"line":13,"column":1,"index":357},"filename":"retry-no-emit.ts"},"fnName":"Foo","memoSlots":0,"memoBlocks":0,"memoValues":0,"prunedMemoBlocks":0,"prunedMemoValues":0}
59 +```
60 +
61 +### Eval output
62 +(kind: ok) {"arr":[1],"arr2":[2]}
63 +{"arr":[2],"arr2":[2]}
64 +logs: [[ 1 ],[ 2 ]]
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/no-emit/retry-no-emit.js new
+19
@@ -0,0 +1,19 @@
1 +// @inferEffectDependencies @noEmit @panicThreshold:"none" @loggerTestOnly
2 +import {print} from 'shared-runtime';
3 +import useEffectWrapper from 'useEffectWrapper';
4 +
5 +function Foo({propVal}) {
6 + const arr = [propVal];
7 + useEffectWrapper(() => print(arr));
8 +
9 + const arr2 = [];
10 + useEffectWrapper(() => arr2.push(propVal));
11 + arr2.push(2);
12 + return {arr, arr2};
13 +}
14 +
15 +export const FIXTURE_ENTRYPOINT = {
16 + fn: Foo,
17 + params: [{propVal: 1}],
18 + sequentialRenders: [{propVal: 1}, {propVal: 2}],
19 +};
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/no-emit/retry-opt-in--no-emit.expect.md new
+60
@@ -0,0 +1,60 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @compilationMode:"all" @inferEffectDependencies @panicThreshold:"none" @noEmit
6 +import {print} from 'shared-runtime';
7 +import useEffectWrapper from 'useEffectWrapper';
8 +
9 +function Foo({propVal}) {
10 + 'use memo';
11 + const arr = [propVal];
12 + useEffectWrapper(() => print(arr));
13 +
14 + const arr2 = [];
15 + useEffectWrapper(() => arr2.push(propVal));
16 + arr2.push(2);
17 +
18 + return {arr, arr2};
19 +}
20 +
21 +export const FIXTURE_ENTRYPOINT = {
22 + fn: Foo,
23 + params: [{propVal: 1}],
24 + sequentialRenders: [{propVal: 1}, {propVal: 2}],
25 +};
26 +
27 +```
28 +
29 +## Code
30 +
31 +```javascript
32 +// @compilationMode:"all" @inferEffectDependencies @panicThreshold:"none" @noEmit
33 +import { print } from "shared-runtime";
34 +import useEffectWrapper from "useEffectWrapper";
35 +
36 +function Foo(t0) {
37 + "use memo";
38 + const { propVal } = t0;
39 +
40 + const arr = [propVal];
41 + useEffectWrapper(() => print(arr), [arr]);
42 +
43 + const arr2 = [];
44 + useEffectWrapper(() => arr2.push(propVal), [arr2, propVal]);
45 + arr2.push(2);
46 + return { arr, arr2 };
47 +}
48 +
49 +export const FIXTURE_ENTRYPOINT = {
50 + fn: Foo,
51 + params: [{ propVal: 1 }],
52 + sequentialRenders: [{ propVal: 1 }, { propVal: 2 }],
53 +};
54 +
55 +```
56 +
57 +### Eval output
58 +(kind: ok) {"arr":[1],"arr2":[2]}
59 +{"arr":[2],"arr2":[2]}
60 +logs: [[ 1 ],[ 2 ]]
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/no-emit/retry-opt-in--no-emit.js new
+21
@@ -0,0 +1,21 @@
1 +// @compilationMode:"all" @inferEffectDependencies @panicThreshold:"none" @noEmit
2 +import {print} from 'shared-runtime';
3 +import useEffectWrapper from 'useEffectWrapper';
4 +
5 +function Foo({propVal}) {
6 + 'use memo';
7 + const arr = [propVal];
8 + useEffectWrapper(() => print(arr));
9 +
10 + const arr2 = [];
11 + useEffectWrapper(() => arr2.push(propVal));
12 + arr2.push(2);
13 +
14 + return {arr, arr2};
15 +}
16 +
17 +export const FIXTURE_ENTRYPOINT = {
18 + fn: Foo,
19 + params: [{propVal: 1}],
20 + sequentialRenders: [{propVal: 1}, {propVal: 2}],
21 +};