[compiler] Run compiler pipeline on 'use no forget'

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

@@ -165,6 +165,12 @@ export type LoggerEvent =
       fnLoc: t.SourceLocation | null;
       detail: Omit<Omit<CompilerErrorDetailOptions, 'severity'>, 'suggestions'>;
     }
+  | {
+      kind: 'CompileSkip';
+      fnLoc: t.SourceLocation | null;
+      reason: string;
+      loc: t.SourceLocation | null;
+    }
   | {
       kind: 'CompileSuccess';
       fnLoc: t.SourceLocation | null;

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

@@ -43,34 +43,23 @@ export type CompilerPass = {
   comments: Array<t.CommentBlock | t.CommentLine>;
   code: string | null;
 };
+const OPT_IN_DIRECTIVES = new Set(['use forget', 'use memo']);
+export const OPT_OUT_DIRECTIVES = new Set(['use no forget', 'use no memo']);
 
 function findDirectiveEnablingMemoization(
   directives: Array<t.Directive>,
-): t.Directive | null {
-  for (const directive of directives) {
-    const directiveValue = directive.value.value;
-    if (directiveValue === 'use forget' || directiveValue === 'use memo') {
-      return directive;
-    }
-  }
-  return null;
+): Array<t.Directive> {
+  return directives.filter(directive =>
+    OPT_IN_DIRECTIVES.has(directive.value.value),
+  );
 }
 
 function findDirectiveDisablingMemoization(
   directives: Array<t.Directive>,
-  options: PluginOptions,
-): t.Directive | null {
-  for (const directive of directives) {
-    const directiveValue = directive.value.value;
-    if (
-      (directiveValue === 'use no forget' ||
-        directiveValue === 'use no memo') &&
-      !options.ignoreUseNoForget
-    ) {
-      return directive;
-    }
-  }
-  return null;
+): Array<t.Directive> {
+  return directives.filter(directive =>
+    OPT_OUT_DIRECTIVES.has(directive.value.value),
+  );
 }
 
 function isCriticalError(err: unknown): boolean {
@@ -102,7 +91,7 @@ export type CompileResult = {
   compiledFn: CodegenFunction;
 };
 
-function handleError(
+function logError(
   err: unknown,
   pass: CompilerPass,
   fnLoc: t.SourceLocation | null,
@@ -131,6 +120,13 @@ function handleError(
       });
     }
   }
+}
+function handleError(
+  err: unknown,
+  pass: CompilerPass,
+  fnLoc: t.SourceLocation | null,
+): void {
+  logError(err, pass, fnLoc);
   if (
     pass.opts.panicThreshold === 'all_errors' ||
     (pass.opts.panicThreshold === 'critical_errors' && isCriticalError(err)) ||
@@ -393,6 +389,17 @@ export function compileProgram(
     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,
+      );
+    }
+
     if (lintError != null) {
       /**
        * Note that Babel does not attach comment nodes to nodes; they are dangling off of the
@@ -404,7 +411,11 @@ export function compileProgram(
         fn,
       );
       if (suppressionsInFunction.length > 0) {
-        handleError(lintError, pass, fn.node.loc ?? null);
+        if (optOutDirectives.length > 0) {
+          logError(lintError, pass, fn.node.loc ?? null);
+        } else {
+          handleError(lintError, pass, fn.node.loc ?? null);
+        }
       }
     }
 
@@ -430,11 +441,50 @@ export function compileProgram(
         prunedMemoValues: compiledFn.prunedMemoValues,
       });
     } catch (err) {
+      /**
+       * If an opt out directive is present, log only instead of throwing and don't mark as
+       * containing a critical error.
+       */
+      if (fn.node.body.type === 'BlockStatement') {
+        if (optOutDirectives.length > 0) {
+          logError(err, pass, fn.node.loc ?? null);
+          return null;
+        }
+      }
       hasCriticalError ||= isCriticalError(err);
       handleError(err, pass, fn.node.loc ?? null);
       return null;
     }
 
+    /**
+     * Always compile functions with opt in directives.
+     */
+    if (optInDirectives.length > 0) {
+      return compiledFn;
+    } else if (pass.opts.compilationMode === 'annotation') {
+      /**
+       * No opt-in directive in annotation mode, so don't insert the compiled function.
+       */
+      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;
+    }
+
     if (!pass.opts.noEmit && !hasCriticalError) {
       return compiledFn;
     }
@@ -481,6 +531,16 @@ export function compileProgram(
     });
   }
 
+  /**
+   * Do not modify source if there is a module scope level opt out directive.
+   */
+  const moduleScopeOptOutDirectives = findDirectiveDisablingMemoization(
+    program.node.directives,
+  );
+  if (moduleScopeOptOutDirectives.length > 0) {
+    return;
+  }
+
   if (pass.opts.gating != null) {
     const error = checkFunctionReferencedBeforeDeclarationAtTopLevel(
       program,
@@ -596,24 +656,6 @@ function shouldSkipCompilation(
     }
   }
 
-  // Top level "use no forget", skip this file entirely
-  const useNoForget = findDirectiveDisablingMemoization(
-    program.node.directives,
-    pass.opts,
-  );
-  if (useNoForget != null) {
-    pass.opts.logger?.logEvent(pass.filename, {
-      kind: 'CompileError',
-      fnLoc: null,
-      detail: {
-        severity: ErrorSeverity.Todo,
-        reason: 'Skipped due to "use no forget" directive.',
-        loc: useNoForget.loc ?? null,
-        suggestions: null,
-      },
-    });
-    return true;
-  }
   const moduleName = pass.opts.runtimeModule ?? 'react/compiler-runtime';
   if (hasMemoCacheFunctionImport(program, moduleName)) {
     return true;
@@ -631,28 +673,8 @@ function getReactFunctionType(
 ): ReactFunctionType | null {
   const hookPattern = environment.hookPattern;
   if (fn.node.body.type === 'BlockStatement') {
-    // Opt-outs disable compilation regardless of mode
-    const useNoForget = findDirectiveDisablingMemoization(
-      fn.node.body.directives,
-      pass.opts,
-    );
-    if (useNoForget != null) {
-      pass.opts.logger?.logEvent(pass.filename, {
-        kind: 'CompileError',
-        fnLoc: fn.node.body.loc ?? null,
-        detail: {
-          severity: ErrorSeverity.Todo,
-          reason: 'Skipped due to "use no forget" directive.',
-          loc: useNoForget.loc ?? null,
-          suggestions: null,
-        },
-      });
-      return null;
-    }
-    // Otherwise opt-ins enable compilation regardless of mode
-    if (findDirectiveEnablingMemoization(fn.node.body.directives) != null) {
+    if (findDirectiveEnablingMemoization(fn.node.body.directives).length > 0)
       return getComponentOrHookLike(fn, hookPattern) ?? 'Other';
-    }
   }
 
   // Component and hook declarations are known components/hooks

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/use-no-forget-with-no-errors.expect.md

@@ -0,0 +1,35 @@
+
+## Input
+
+```javascript
+function Component() {
+  'use no forget';
+  return <div>Hello World</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+  isComponent: true,
+};
+
+```
+
+## Code
+
+```javascript
+function Component() {
+  "use no forget";
+  return <div>Hello World</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+  isComponent: true,
+};
+
+```
+
+### Eval output
+(kind: ok) <div>Hello World</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/use-no-forget-with-no-errors.js

@@ -0,0 +1,10 @@
+function Component() {
+  'use no forget';
+  return <div>Hello World</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+  isComponent: true,
+};

compiler/packages/babel-plugin-react-compiler/src/index.ts

@@ -18,6 +18,7 @@ export {
   compileProgram,
   parsePluginOptions,
   run,
+  OPT_OUT_DIRECTIVES,
   type CompilerPipelineValue,
   type PluginOptions,
 } from './Entrypoint';