[compiler][be] Make program traversal more readable

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)

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Imports.ts

@@ -18,8 +18,9 @@ import {
 import {getOrInsertWith} from '../Utils/utils';
 import {ExternalFunction, isHookName} from '../HIR/Environment';
 import {Err, Ok, Result} from '../Utils/Result';
-import {CompilerReactTarget} from './Options';
+import {LoggerEvent, PluginOptions} from './Options';
-import {getReactCompilerRuntimeModule} from './Program';
+import {BabelFn, getReactCompilerRuntimeModule} from './Program';
+import {SuppressionRange} from './Suppression';
 
 export function validateRestrictedImports(
   path: NodePath<t.Program>,
@@ -52,32 +53,61 @@ export function validateRestrictedImports(
   }
 }
 
+type ProgramContextOptions = {
+  program: NodePath<t.Program>;
+  suppressions: Array<SuppressionRange>;
+  opts: PluginOptions;
+  filename: string | null;
+  code: string | null;
+};
 export class ProgramContext {
-  /* Program and environment context */
+  /**
+   * Program and environment context
+   */
   scope: BabelScope;
+  opts: PluginOptions;
+  filename: string | null;
+  code: string | null;
   reactRuntimeModule: string;
-  hookPattern: string | null;
+  suppressions: Array<SuppressionRange>;
 
+  /*
+   * This is a hack to work around what seems to be a Babel bug. Babel doesn't
+   * consistently respect the `skip()` function to avoid revisiting a node within
+   * a pass, so we use this set to track nodes that we have compiled.
+   */
+  alreadyCompiled: WeakSet<object> | Set<object> = new (WeakSet ?? Set)();
   // known generated or referenced identifiers in the program
   knownReferencedNames: Set<string> = new Set();
   // generated imports
   imports: Map<string, Map<string, NonLocalImportSpecifier>> = new Map();
 
-  constructor(
+  /**
-    program: NodePath<t.Program>,
+   * Metadata from compilation
-    reactRuntimeModule: CompilerReactTarget,
+   */
-    hookPattern: string | null,
+  retryErrors: Array<{fn: BabelFn; error: CompilerError}> = [];
-  ) {
+  inferredEffectLocations: Set<t.SourceLocation> = new Set();
-    this.hookPattern = hookPattern;
+
+  constructor({
+    program,
+    suppressions,
+    opts,
+    filename,
+    code,
+  }: ProgramContextOptions) {
     this.scope = program.scope;
-    this.reactRuntimeModule = getReactCompilerRuntimeModule(reactRuntimeModule);
+    this.opts = opts;
+    this.filename = filename;
+    this.code = code;
+    this.reactRuntimeModule = getReactCompilerRuntimeModule(opts.target);
+    this.suppressions = suppressions;
   }
 
   isHookName(name: string): boolean {
-    if (this.hookPattern == null) {
+    if (this.opts.environment.hookPattern == null) {
       return isHookName(name);
     } else {
-      const match = new RegExp(this.hookPattern).exec(name);
+      const match = new RegExp(this.opts.environment.hookPattern).exec(name);
       return (
         match != null && typeof match[1] === 'string' && isHookName(match[1])
       );
@@ -179,6 +209,12 @@ export class ProgramContext {
     });
     return Err(error);
   }
+
+  logEvent(event: LoggerEvent): void {
+    if (this.opts.logger != null) {
+      this.opts.logger.logEvent(this.filename, event);
+    }
+  }
 }
 
 function getExistingImports(

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts

@@ -12,7 +12,7 @@ import {
   CompilerErrorDetail,
   ErrorSeverity,
 } from '../CompilerError';
-import {EnvironmentConfig, ReactFunctionType} from '../HIR/Environment';
+import {ReactFunctionType} from '../HIR/Environment';
 import {CodegenFunction} from '../ReactiveScopes';
 import {isComponentDeclaration} from '../Utils/ComponentDeclaration';
 import {isHookDeclaration} from '../Utils/HookDeclaration';
@@ -43,17 +43,21 @@ export const OPT_OUT_DIRECTIVES = new Set(['use no forget', 'use no memo']);
 
 export function findDirectiveEnablingMemoization(
   directives: Array<t.Directive>,
-): Array<t.Directive> {
+): t.Directive | null {
-  return directives.filter(directive =>
+  return (
-    OPT_IN_DIRECTIVES.has(directive.value.value),
+    directives.find(directive =>
+      OPT_IN_DIRECTIVES.has(directive.value.value),
+    ) ?? null
   );
 }
 
 export function findDirectiveDisablingMemoization(
   directives: Array<t.Directive>,
-): Array<t.Directive> {
+): t.Directive | null {
-  return directives.filter(directive =>
+  return (
-    OPT_OUT_DIRECTIVES.has(directive.value.value),
+    directives.find(directive =>
+      OPT_OUT_DIRECTIVES.has(directive.value.value),
+    ) ?? null
   );
 }
 
@@ -88,13 +92,16 @@ export type CompileResult = {
 
 function logError(
   err: unknown,
-  pass: CompilerPass,
+  context: {
+    opts: PluginOptions;
+    filename: string | null;
+  },
   fnLoc: t.SourceLocation | null,
 ): void {
-  if (pass.opts.logger) {
+  if (context.opts.logger) {
     if (err instanceof CompilerError) {
       for (const detail of err.details) {
-        pass.opts.logger.logEvent(pass.filename, {
+        context.opts.logger.logEvent(context.filename, {
           kind: 'CompileError',
           fnLoc,
           detail: detail.options,
@@ -108,7 +115,7 @@ function logError(
         stringifiedError = err?.toString() ?? '[ null ]';
       }
 
-      pass.opts.logger.logEvent(pass.filename, {
+      context.opts.logger.logEvent(context.filename, {
         kind: 'PipelineError',
         fnLoc,
         data: stringifiedError,
@@ -118,13 +125,17 @@ function logError(
 }
 function handleError(
   err: unknown,
-  pass: CompilerPass,
+  context: {
+    opts: PluginOptions;
+    filename: string | null;
+  },
   fnLoc: t.SourceLocation | null,
 ): void {
-  logError(err, pass, fnLoc);
+  logError(err, context, fnLoc);
   if (
-    pass.opts.panicThreshold === 'all_errors' ||
+    context.opts.panicThreshold === 'all_errors' ||
-    (pass.opts.panicThreshold === 'critical_errors' && isCriticalError(err)) ||
+    (context.opts.panicThreshold === 'critical_errors' &&
+      isCriticalError(err)) ||
     isConfigError(err) // Always throws regardless of panic threshold
   ) {
     throw err;
@@ -187,7 +198,6 @@ export function createNewFunctionNode(
     }
   }
   // Avoid visiting the new transformed version
-  ALREADY_COMPILED.add(transformedFn);
   return transformedFn;
 }
 
@@ -239,13 +249,6 @@ function insertNewOutlinedFunctionNode(
   }
 }
 
-/*
- * This is a hack to work around what seems to be a Babel bug. Babel doesn't
- * consistently respect the `skip()` function to avoid revisiting a node within
- * a pass, so we use this set to track nodes that we have compiled.
- */
-const ALREADY_COMPILED: WeakSet<object> | Set<object> = new (WeakSet ?? Set)();
-
 const DEFAULT_ESLINT_SUPPRESSIONS = [
   'react-hooks/exhaustive-deps',
   'react-hooks/rules-of-hooks',
@@ -268,41 +271,43 @@ function isFilePartOfSources(
   return false;
 }
 
-export type CompileProgramResult = {
+export type CompileProgramMetadata = {
   retryErrors: Array<{fn: BabelFn; error: CompilerError}>;
   inferredEffectLocations: Set<t.SourceLocation>;
 };
 /**
- * `compileProgram` is directly invoked by the react-compiler babel plugin, so
+ * Main entrypoint for React Compiler.
- * exceptions thrown by this function will fail the babel build.
+ *
- * - call `handleError` if your error is recoverable.
+ * @param program The Babel program node to compile
- *   Unless the error is a warning / info diagnostic, compilation of a function
+ * @param pass Compiler configuration and context
- *   / entire file should also be skipped.
+ * @returns Compilation results or null if compilation was skipped
- * - throw an exception if the error is fatal / not recoverable.
- *   Examples of this are invalid compiler configs or failure to codegen outlined
- *   functions *after* already emitting optimized components / hooks that invoke
- *   the outlined functions.
  */
 export function compileProgram(
   program: NodePath<t.Program>,
   pass: CompilerPass,
-): CompileProgramResult | null {
+): CompileProgramMetadata | null {
+  /**
+   * This is directly invoked by the react-compiler babel plugin, so exceptions
+   * thrown by this function will fail the babel build.
+   * - call `handleError` if your error is recoverable.
+   *   Unless the error is a warning / info diagnostic, compilation of a function
+   *   / entire file should also be skipped.
+   * - throw an exception if the error is fatal / not recoverable.
+   *   Examples of this are invalid compiler configs or failure to codegen outlined
+   *   functions *after* already emitting optimized components / hooks that invoke
+   *   the outlined functions.
+   */
   if (shouldSkipCompilation(program, pass)) {
     return null;
   }
-
+  const restrictedImportsErr = validateRestrictedImports(
-  const environment = pass.opts.environment;
+    program,
-  const restrictedImportsErr = validateRestrictedImports(program, environment);
+    pass.opts.environment,
+  );
   if (restrictedImportsErr) {
     handleError(restrictedImportsErr, pass, null);
     return null;
   }
-
-  const programContext = new ProgramContext(
-    program,
-    pass.opts.target,
-    environment.hookPattern,
-  );
   /*
    * Record lint errors and critical errors as depending on Forget's config,
    * we may still need to run Forget's analysis on every function (even if we
@@ -313,16 +318,88 @@ export function compileProgram(
     pass.opts.eslintSuppressionRules ?? DEFAULT_ESLINT_SUPPRESSIONS,
     pass.opts.flowSuppressions,
   );
-  const queue: Array<{
+
-    kind: 'original' | 'outlined';
+  const programContext = new ProgramContext({
-    fn: BabelFn;
+    program: program,
-    fnType: ReactFunctionType;
+    opts: pass.opts,
-  }> = [];
+    filename: pass.filename,
+    code: pass.code,
+    suppressions,
+  });
+
+  const queue: Array<CompileSource> = findFunctionsToCompile(
+    program,
+    pass,
+    programContext,
+  );
   const compiledFns: Array<CompileResult> = [];
 
+  while (queue.length !== 0) {
+    const current = queue.shift()!;
+    const compiled = processFn(current.fn, current.fnType, programContext);
+
+    if (compiled != null) {
+      for (const outlined of compiled.outlined) {
+        CompilerError.invariant(outlined.fn.outlined.length === 0, {
+          reason: 'Unexpected nested outlined functions',
+          loc: outlined.fn.loc,
+        });
+        const fn = insertNewOutlinedFunctionNode(
+          program,
+          current.fn,
+          outlined.fn,
+        );
+        fn.skip();
+        programContext.alreadyCompiled.add(fn.node);
+        if (outlined.type !== null) {
+          queue.push({
+            kind: 'outlined',
+            fn,
+            fnType: outlined.type,
+          });
+        }
+      }
+      compiledFns.push({
+        kind: current.kind,
+        originalFn: current.fn,
+        compiledFn: compiled,
+      });
+    }
+  }
+
+  // Avoid modifying the program if we find a program level opt-out
+  if (findDirectiveDisablingMemoization(program.node.directives) != null) {
+    return null;
+  }
+
+  // Insert React Compiler generated functions into the Babel AST
+  applyCompiledFunctions(program, compiledFns, pass, programContext);
+
+  return {
+    retryErrors: programContext.retryErrors,
+    inferredEffectLocations: programContext.inferredEffectLocations,
+  };
+}
+
+type CompileSource = {
+  kind: 'original' | 'outlined';
+  fn: BabelFn;
+  fnType: ReactFunctionType;
+};
+/**
+ * Find all React components and hooks that need to be compiled
+ *
+ * @returns An array of React functions from @param program to transform
+ */
+function findFunctionsToCompile(
+  program: NodePath<t.Program>,
+  pass: CompilerPass,
+  programContext: ProgramContext,
+): Array<CompileSource> {
+  const queue: Array<CompileSource> = [];
   const traverseFunction = (fn: BabelFn, pass: CompilerPass): void => {
-    const fnType = getReactFunctionType(fn, pass, environment);
+    const fnType = getReactFunctionType(fn, pass);
-    if (fnType === null || ALREADY_COMPILED.has(fn.node)) {
+    if (fnType === null || programContext.alreadyCompiled.has(fn.node)) {
       return;
     }
 
@@ -331,7 +408,7 @@ export function compileProgram(
      * traversal will loop infinitely.
      * Ensure we avoid visiting the original function again.
      */
-    ALREADY_COMPILED.add(fn.node);
+    programContext.alreadyCompiled.add(fn.node);
     fn.skip();
 
     queue.push({kind: 'original', fn, fnType});
@@ -346,7 +423,6 @@ export function compileProgram(
          * can reference `this` which is unsafe for compilation
          */
         node.skip();
-        return;
       },
 
       ClassExpression(node: NodePath<t.ClassExpression>) {
@@ -355,7 +431,6 @@ export function compileProgram(
          * can reference `this` which is unsafe for compilation
          */
         node.skip();
-        return;
       },
 
       FunctionDeclaration: traverseFunction,
@@ -370,205 +445,205 @@ export function compileProgram(
       filename: pass.filename ?? null,
     },
   );
-  const retryErrors: Array<{fn: BabelFn; error: CompilerError}> = [];
+  return queue;
-  const inferredEffectLocations = new Set<t.SourceLocation>();
+}
-  const processFn = (
-    fn: BabelFn,
-    fnType: ReactFunctionType,
-  ): null | CodegenFunction => {
-    let optInDirectives: Array<t.Directive> = [];
-    let optOutDirectives: Array<t.Directive> = [];
-    if (fn.node.body.type === 'BlockStatement') {
-      optInDirectives = findDirectiveEnablingMemoization(
-        fn.node.body.directives,
-      );
-      optOutDirectives = findDirectiveDisablingMemoization(
-        fn.node.body.directives,
-      );
-    }
 
-    /**
+/**
-     * Note that Babel does not attach comment nodes to nodes; they are dangling off of the
+ * Try to compile a source function, taking into account all local suppressions,
-     * Program node itself. We need to figure out whether an eslint suppression range
+ * opt-ins, and opt-outs.
-     * applies to this function first.
+ *
-     */
+ * Errors encountered during compilation are either logged (if recoverable) or
-    const suppressionsInFunction = filterSuppressionsThatAffectFunction(
+ * thrown (if non-recoverable).
-      suppressions,
+ *
-      fn,
+ * @returns the compiled function or null if the function was skipped (due to
-    );
+ * config settings and/or outputs)
-    let compileResult:
+ */
-      | {kind: 'compile'; compiledFn: CodegenFunction}
+function processFn(
-      | {kind: 'error'; error: unknown};
+  fn: BabelFn,
-    if (suppressionsInFunction.length > 0) {
+  fnType: ReactFunctionType,
-      compileResult = {
+  programContext: ProgramContext,
-        kind: 'error',
+): null | CodegenFunction {
-        error: suppressionsToCompilerError(suppressionsInFunction),
+  let directives;
-      };
+  if (fn.node.body.type !== 'BlockStatement') {
+    directives = {optIn: null, optOut: null};
+  } else {
+    directives = {
+      optIn: findDirectiveEnablingMemoization(fn.node.body.directives),
+      optOut: findDirectiveDisablingMemoization(fn.node.body.directives),
+    };
+  }
+
+  let compiledFn: CodegenFunction;
+  const compileResult = tryCompileFunction(fn, fnType, programContext);
+  if (compileResult.kind === 'error') {
+    if (directives.optOut != null) {
+      logError(compileResult.error, programContext, fn.node.loc ?? null);
     } else {
-      try {
+      handleError(compileResult.error, programContext, fn.node.loc ?? null);
-        compileResult = {
-          kind: 'compile',
-          compiledFn: compileFn(
-            fn,
-            environment,
-            fnType,
-            'all_features',
-            programContext,
-            pass.opts.logger,
-            pass.filename,
-            pass.code,
-          ),
-        };
-      } catch (err) {
-        compileResult = {kind: 'error', error: err};
-      }
     }
-
+    const retryResult = retryCompileFunction(fn, fnType, programContext);
-    if (compileResult.kind === 'error') {
+    if (retryResult == null) {
-      /**
-       * If an opt out directive is present, log only instead of throwing and don't mark as
-       * containing a critical error.
-       */
-      if (optOutDirectives.length > 0) {
-        logError(compileResult.error, pass, fn.node.loc ?? null);
-      } else {
-        handleError(compileResult.error, pass, fn.node.loc ?? null);
-      }
-      // If non-memoization features are enabled, retry regardless of error kind
-      if (
-        !(environment.enableFire || environment.inferEffectDependencies != null)
-      ) {
-        return null;
-      }
-      try {
-        compileResult = {
-          kind: 'compile',
-          compiledFn: compileFn(
-            fn,
-            environment,
-            fnType,
-            'no_inferred_memo',
-            programContext,
-            pass.opts.logger,
-            pass.filename,
-            pass.code,
-          ),
-        };
-        if (
-          !compileResult.compiledFn.hasFireRewrite &&
-          !compileResult.compiledFn.hasInferredEffect
-        ) {
-          return null;
-        }
-      } catch (err) {
-        // TODO: we might want to log error here, but this will also result in duplicate logging
-        if (err instanceof CompilerError) {
-          retryErrors.push({fn, error: err});
-        }
-        return null;
-      }
-    }
-
-    /**
-     * Otherwise if 'use no forget/memo' is present, we still run the code through the compiler
-     * for validation but we don't mutate the babel AST. This allows us to flag if there is an
-     * unused 'use no forget/memo' directive.
-     */
-    if (pass.opts.ignoreUseNoForget === false && optOutDirectives.length > 0) {
-      for (const directive of optOutDirectives) {
-        pass.opts.logger?.logEvent(pass.filename, {
-          kind: 'CompileSkip',
-          fnLoc: fn.node.body.loc ?? null,
-          reason: `Skipped due to '${directive.value.value}' directive.`,
-          loc: directive.loc ?? null,
-        });
-      }
       return null;
     }
+    compiledFn = retryResult;
+  } else {
+    compiledFn = compileResult.compiledFn;
+  }
 
-    pass.opts.logger?.logEvent(pass.filename, {
+  /**
-      kind: 'CompileSuccess',
+   * Otherwise if 'use no forget/memo' is present, we still run the code through the compiler
-      fnLoc: fn.node.loc ?? null,
+   * for validation but we don't mutate the babel AST. This allows us to flag if there is an
-      fnName: compileResult.compiledFn.id?.name ?? null,
+   * unused 'use no forget/memo' directive.
-      memoSlots: compileResult.compiledFn.memoSlotsUsed,
+   */
-      memoBlocks: compileResult.compiledFn.memoBlocks,
+  if (
-      memoValues: compileResult.compiledFn.memoValues,
+    programContext.opts.ignoreUseNoForget === false &&
-      prunedMemoBlocks: compileResult.compiledFn.prunedMemoBlocks,
+    directives.optOut != null
-      prunedMemoValues: compileResult.compiledFn.prunedMemoValues,
+  ) {
+    programContext.logEvent({
+      kind: 'CompileSkip',
+      fnLoc: fn.node.body.loc ?? null,
+      reason: `Skipped due to '${directives.optOut.value}' directive.`,
+      loc: directives.optOut.loc ?? null,
     });
+    return null;
+  }
+  programContext.logEvent({
+    kind: 'CompileSuccess',
+    fnLoc: fn.node.loc ?? null,
+    fnName: compiledFn.id?.name ?? null,
+    memoSlots: compiledFn.memoSlotsUsed,
+    memoBlocks: compiledFn.memoBlocks,
+    memoValues: compiledFn.memoValues,
+    prunedMemoBlocks: compiledFn.prunedMemoBlocks,
+    prunedMemoValues: compiledFn.prunedMemoValues,
+  });
 
+  /**
+   * Always compile functions with opt in directives.
+   */
+  if (directives.optIn != null) {
+    return compiledFn;
+  } else if (programContext.opts.compilationMode === 'annotation') {
     /**
-     * Always compile functions with opt in directives.
+     * If no opt-in directive is found and the compiler is configured in
+     * annotation mode, don't insert the compiled function.
      */
-    if (optInDirectives.length > 0) {
+    return null;
-      return compileResult.compiledFn;
+  } else if (programContext.opts.noEmit) {
-    } else if (pass.opts.compilationMode === 'annotation') {
-      /**
-       * No opt-in directive in annotation mode, so don't insert the compiled function.
-       */
-      return null;
-    }
-
-    if (!pass.opts.noEmit) {
-      return compileResult.compiledFn;
-    }
     /**
      * inferEffectDependencies + noEmit is currently only used for linting. In
      * this mode, add source locations for where the compiler *can* infer effect
      * dependencies.
      */
-    for (const loc of compileResult.compiledFn.inferredEffectLocations) {
+    for (const loc of compiledFn.inferredEffectLocations) {
-      if (loc !== GeneratedSource) inferredEffectLocations.add(loc);
+      if (loc !== GeneratedSource) {
-    }
+        programContext.inferredEffectLocations.add(loc);
-    return null;
-  };
-
-  while (queue.length !== 0) {
-    const current = queue.shift()!;
-    const compiled = processFn(current.fn, current.fnType);
-    if (compiled === null) {
-      continue;
-    }
-    for (const outlined of compiled.outlined) {
-      CompilerError.invariant(outlined.fn.outlined.length === 0, {
-        reason: 'Unexpected nested outlined functions',
-        loc: outlined.fn.loc,
-      });
-      const fn = insertNewOutlinedFunctionNode(
-        program,
-        current.fn,
-        outlined.fn,
-      );
-      fn.skip();
-      ALREADY_COMPILED.add(fn.node);
-      if (outlined.type !== null) {
-        queue.push({
-          kind: 'outlined',
-          fn,
-          fnType: outlined.type,
-        });
       }
     }
-    compiledFns.push({
+    return null;
-      kind: current.kind,
+  } else {
-      compiledFn: compiled,
+    return compiledFn;
-      originalFn: current.fn,
+  }
-    });
+}
+
+function tryCompileFunction(
+  fn: BabelFn,
+  fnType: ReactFunctionType,
+  programContext: ProgramContext,
+):
+  | {kind: 'compile'; compiledFn: CodegenFunction}
+  | {kind: 'error'; error: unknown} {
+  /**
+   * Note that Babel does not attach comment nodes to nodes; they are dangling off of the
+   * Program node itself. We need to figure out whether an eslint suppression range
+   * applies to this function first.
+   */
+  const suppressionsInFunction = filterSuppressionsThatAffectFunction(
+    programContext.suppressions,
+    fn,
+  );
+  if (suppressionsInFunction.length > 0) {
+    return {
+      kind: 'error',
+      error: suppressionsToCompilerError(suppressionsInFunction),
+    };
   }
 
-  /**
+  try {
-   * Do not modify source if there is a module scope level opt out directive.
+    return {
-   */
+      kind: 'compile',
-  const moduleScopeOptOutDirectives = findDirectiveDisablingMemoization(
+      compiledFn: compileFn(
-    program.node.directives,
+        fn,
-  );
+        programContext.opts.environment,
-  if (moduleScopeOptOutDirectives.length > 0) {
+        fnType,
+        'all_features',
+        programContext,
+        programContext.opts.logger,
+        programContext.filename,
+        programContext.code,
+      ),
+    };
+  } catch (err) {
+    return {kind: 'error', error: err};
+  }
+}
+
+/**
+ * If non-memo feature flags are enabled, retry compilation with a more minimal
+ * feature set.
+ *
+ * @returns a CodegenFunction if retry was successful
+ */
+function retryCompileFunction(
+  fn: BabelFn,
+  fnType: ReactFunctionType,
+  programContext: ProgramContext,
+): CodegenFunction | null {
+  const environment = programContext.opts.environment;
+  if (
+    !(environment.enableFire || environment.inferEffectDependencies != null)
+  ) {
     return null;
   }
-  /*
+  /**
-   * Only insert Forget-ified functions if we have not encountered a critical
+   * Note that function suppressions are not checked in the retry pipeline, as
-   * error elsewhere in the file, regardless of bailout mode.
+   * they only affect auto-memoization features.
    */
+  try {
+    const retryResult = compileFn(
+      fn,
+      environment,
+      fnType,
+      'no_inferred_memo',
+      programContext,
+      programContext.opts.logger,
+      programContext.filename,
+      programContext.code,
+    );
+
+    if (!retryResult.hasFireRewrite && !retryResult.hasInferredEffect) {
+      return null;
+    }
+    return retryResult;
+  } catch (err) {
+    // TODO: we might want to log error here, but this will also result in duplicate logging
+    if (err instanceof CompilerError) {
+      programContext.retryErrors.push({fn, error: err});
+    }
+    return null;
+  }
+}
+
+/**
+ * Applies React Compiler generated functions to the babel AST by replacing
+ * existing functions in place or inserting new declarations.
+ */
+function applyCompiledFunctions(
+  program: NodePath<t.Program>,
+  compiledFns: Array<CompileResult>,
+  pass: CompilerPass,
+  programContext: ProgramContext,
+): void {
   const referencedBeforeDeclared =
     pass.opts.gating != null
       ? getFunctionReferencedBeforeDeclarationAtTopLevel(program, compiledFns)
@@ -576,6 +651,7 @@ export function compileProgram(
   for (const result of compiledFns) {
     const {kind, originalFn, compiledFn} = result;
     const transformedFn = createNewFunctionNode(originalFn, compiledFn);
+    programContext.alreadyCompiled.add(transformedFn);
 
     if (referencedBeforeDeclared != null && kind === 'original') {
       CompilerError.invariant(pass.opts.gating != null, {
@@ -598,7 +674,6 @@ export function compileProgram(
   if (compiledFns.length > 0) {
     addImportsToProgram(program, programContext);
   }
-  return {retryErrors, inferredEffectLocations};
 }
 
 function shouldSkipCompilation(
@@ -640,14 +715,10 @@ function shouldSkipCompilation(
 function getReactFunctionType(
   fn: BabelFn,
   pass: CompilerPass,
-  /**
-   * TODO(mofeiZ): remove once we validate PluginOptions with Zod
-   */
-  environment: EnvironmentConfig,
 ): ReactFunctionType | null {
-  const hookPattern = environment.hookPattern;
+  const hookPattern = pass.opts.environment.hookPattern;
   if (fn.node.body.type === 'BlockStatement') {
-    if (findDirectiveEnablingMemoization(fn.node.body.directives).length > 0)
+    if (findDirectiveEnablingMemoization(fn.node.body.directives) != null)
       return getComponentOrHookLike(fn, hookPattern) ?? 'Other';
   }
 

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/ValidateNoUntransformedReferences.ts

@@ -18,7 +18,7 @@ import {
 import {getOrInsertWith} from '../Utils/utils';
 import {Environment} from '../HIR';
 import {DEFAULT_EXPORT} from '../HIR/Environment';
-import {CompileProgramResult} from './Program';
+import {CompileProgramMetadata} from './Program';
 
 function throwInvalidReact(
   options: Omit<CompilerErrorDetailOptions, 'severity'>,
@@ -109,7 +109,7 @@ export default function validateNoUntransformedReferences(
   filename: string | null,
   logger: Logger | null,
   env: EnvironmentConfig,
-  compileResult: CompileProgramResult | null,
+  compileResult: CompileProgramMetadata | null,
 ): void {
   const moduleLoadChecks = new Map<
     string,
@@ -236,7 +236,7 @@ function transformProgram(
   moduleLoadChecks: Map<string, Map<string, CheckInvalidReferenceFn>>,
   filename: string | null,
   logger: Logger | null,
-  compileResult: CompileProgramResult | null,
+  compileResult: CompileProgramMetadata | null,
 ): void {
   const traversalState: TraversalState = {
     shouldInvalidateScopes: true,

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/no-emit/retry-no-emit.expect.md

@@ -0,0 +1,64 @@
+
+## Input
+
+```javascript
+// @inferEffectDependencies @noEmit @panicThreshold:"none" @loggerTestOnly
+import {print} from 'shared-runtime';
+import useEffectWrapper from 'useEffectWrapper';
+
+function Foo({propVal}) {
+  const arr = [propVal];
+  useEffectWrapper(() => print(arr));
+
+  const arr2 = [];
+  useEffectWrapper(() => arr2.push(propVal));
+  arr2.push(2);
+  return {arr, arr2};
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{propVal: 1}],
+  sequentialRenders: [{propVal: 1}, {propVal: 2}],
+};
+
+```
+
+## Code
+
+```javascript
+// @inferEffectDependencies @noEmit @panicThreshold:"none" @loggerTestOnly
+import { print } from "shared-runtime";
+import useEffectWrapper from "useEffectWrapper";
+
+function Foo({ propVal }) {
+  const arr = [propVal];
+  useEffectWrapper(() => print(arr));
+
+  const arr2 = [];
+  useEffectWrapper(() => arr2.push(propVal));
+  arr2.push(2);
+  return { arr, arr2 };
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{ propVal: 1 }],
+  sequentialRenders: [{ propVal: 1 }, { propVal: 2 }],
+};
+
+```
+
+## Logs
+
+```
+{"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"}}
+{"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"}]}
+{"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"}]}
+{"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}
+```
+      
+### Eval output
+(kind: ok) {"arr":[1],"arr2":[2]}
+{"arr":[2],"arr2":[2]}
+logs: [[ 1 ],[ 2 ]]

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/no-emit/retry-no-emit.js

@@ -0,0 +1,19 @@
+// @inferEffectDependencies @noEmit @panicThreshold:"none" @loggerTestOnly
+import {print} from 'shared-runtime';
+import useEffectWrapper from 'useEffectWrapper';
+
+function Foo({propVal}) {
+  const arr = [propVal];
+  useEffectWrapper(() => print(arr));
+
+  const arr2 = [];
+  useEffectWrapper(() => arr2.push(propVal));
+  arr2.push(2);
+  return {arr, arr2};
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{propVal: 1}],
+  sequentialRenders: [{propVal: 1}, {propVal: 2}],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/no-emit/retry-opt-in--no-emit.expect.md

@@ -0,0 +1,60 @@
+
+## Input
+
+```javascript
+// @compilationMode:"all" @inferEffectDependencies @panicThreshold:"none" @noEmit
+import {print} from 'shared-runtime';
+import useEffectWrapper from 'useEffectWrapper';
+
+function Foo({propVal}) {
+  'use memo';
+  const arr = [propVal];
+  useEffectWrapper(() => print(arr));
+
+  const arr2 = [];
+  useEffectWrapper(() => arr2.push(propVal));
+  arr2.push(2);
+
+  return {arr, arr2};
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{propVal: 1}],
+  sequentialRenders: [{propVal: 1}, {propVal: 2}],
+};
+
+```
+
+## Code
+
+```javascript
+// @compilationMode:"all" @inferEffectDependencies @panicThreshold:"none" @noEmit
+import { print } from "shared-runtime";
+import useEffectWrapper from "useEffectWrapper";
+
+function Foo(t0) {
+  "use memo";
+  const { propVal } = t0;
+
+  const arr = [propVal];
+  useEffectWrapper(() => print(arr), [arr]);
+
+  const arr2 = [];
+  useEffectWrapper(() => arr2.push(propVal), [arr2, propVal]);
+  arr2.push(2);
+  return { arr, arr2 };
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{ propVal: 1 }],
+  sequentialRenders: [{ propVal: 1 }, { propVal: 2 }],
+};
+
+```
+      
+### Eval output
+(kind: ok) {"arr":[1],"arr2":[2]}
+{"arr":[2],"arr2":[2]}
+logs: [[ 1 ],[ 2 ]]

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/no-emit/retry-opt-in--no-emit.js

@@ -0,0 +1,21 @@
+// @compilationMode:"all" @inferEffectDependencies @panicThreshold:"none" @noEmit
+import {print} from 'shared-runtime';
+import useEffectWrapper from 'useEffectWrapper';
+
+function Foo({propVal}) {
+  'use memo';
+  const arr = [propVal];
+  useEffectWrapper(() => print(arr));
+
+  const arr2 = [];
+  useEffectWrapper(() => arr2.push(propVal));
+  arr2.push(2);
+
+  return {arr, arr2};
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{propVal: 1}],
+  sequentialRenders: [{propVal: 1}, {propVal: 2}],
+};