[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)) ||
@@ -430,11 +426,65 @@ 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') {
+        const optOutDirectives = findDirectiveDisablingMemoization(
+          fn.node.body.directives,
+        );
+        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;
     }
 
+    if (fn.node.body.type === 'BlockStatement') {
+      const optInDirectives = findDirectiveEnablingMemoization(
+        fn.node.body.directives,
+      );
+      const optOutDirectives = findDirectiveDisablingMemoization(
+        fn.node.body.directives,
+      );
+
+      /**
+       * 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/error.use-no-forget-multiple-with-eslint-suppression.expect.md

@@ -0,0 +1,38 @@
+
+## Input
+
+```javascript
+import {useRef} from 'react';
+
+const useControllableState = options => {};
+function NoopComponent() {}
+
+function Component() {
+  'use no forget';
+  const ref = useRef(null);
+  // eslint-disable-next-line react-hooks/rules-of-hooks
+  ref.current = 'bad';
+  return <button ref={ref} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+};
+
+```
+
+
+## Error
+
+```
+   7 |   'use no forget';
+   8 |   const ref = useRef(null);
+>  9 |   // eslint-disable-next-line react-hooks/rules-of-hooks
+     |   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ InvalidReact: React Compiler has skipped optimizing this component because one or more React ESLint rules were disabled. React Compiler only works when your components follow all the rules of React, disabling them may result in unexpected or incorrect behavior. eslint-disable-next-line react-hooks/rules-of-hooks (9:9)
+  10 |   ref.current = 'bad';
+  11 |   return <button ref={ref} />;
+  12 | }
+```
+
+

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

@@ -0,0 +1,35 @@
+
+## Input
+
+```javascript
+import {useRef} from 'react';
+
+function Component() {
+  'use no forget';
+  const ref = useRef(null);
+  // eslint-disable-next-line react-hooks/rules-of-hooks
+  ref.current = 'bad';
+  return <button ref={ref} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+};
+
+```
+
+
+## Error
+
+```
+  4 |   'use no forget';
+  5 |   const ref = useRef(null);
+> 6 |   // eslint-disable-next-line react-hooks/rules-of-hooks
+    |   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ InvalidReact: React Compiler has skipped optimizing this component because one or more React ESLint rules were disabled. React Compiler only works when your components follow all the rules of React, disabling them may result in unexpected or incorrect behavior. eslint-disable-next-line react-hooks/rules-of-hooks (6:6)
+  7 |   ref.current = 'bad';
+  8 |   return <button ref={ref} />;
+  9 | }
+```
+
+