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

This PR updates the babel plugin to continue the compilation pipeline as
normal on components/hooks that have been opted out using a directive.
Instead, we no longer emit the compiled function when the directive is
present.

Previously, we would skip over the entire pipeline. By continuing to
enter the pipeline, we'll be able to detect if there are unused
directives.

The end result is:

- (no change) 'use forget' will always opt into compilation
- (new) 'use no forget' will opt out of compilation but continue to log
  errors without throwing them. This means that a Program containing
multiple functions (some of which are opted out) will continue to
compile correctly

ghstack-source-id: 895dc05
Pull Request resolved: #30720

Author: poteto

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,35 @@ 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;
-    }
+  node: NodePath<t.Program> | NodePath<t.BlockStatement>,
+): Array<NodePath<t.Directive>> {
+  const directives = node.get('directives');
+  if (Array.isArray(directives)) {
+    return (directives as Array<NodePath<t.Directive>>).filter(
+      directive =>
+        directive.isDirective() &&
+        OPT_IN_DIRECTIVES.has(directive.node.value.value),
+    );
   }
-  return null;
+  return [];
 }
 
 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;
-    }
+  node: NodePath<t.Program> | NodePath<t.BlockStatement>,
+): Array<NodePath<t.Directive>> {
+  const directives = node.get('directives');
+  if (Array.isArray(directives)) {
+    return (directives as Array<NodePath<t.Directive>>).filter(
+      directive =>
+        directive.isDirective() &&
+        OPT_OUT_DIRECTIVES.has(directive.node.value.value),
+    );
   }
-  return null;
+  return [];
 }
 
 function isCriticalError(err: unknown): boolean {
@@ -102,7 +103,7 @@ export type CompileResult = {
   compiledFn: CodegenFunction;
 };
 
-function handleError(
+function logError(
   err: unknown,
   pass: CompilerPass,
   fnLoc: t.SourceLocation | null,
@@ -131,6 +132,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)) ||
@@ -408,6 +416,13 @@ export function compileProgram(
       }
     }
 
+    // TS doesn't seem to be able to resolve this type correctly
+    const body = fn.get('body');
+    CompilerError.invariant(Array.isArray(body) === false, {
+      reason: 'Expected fn to only have one body',
+      description: `Got ${body}`,
+      loc: fn.node.loc ?? null,
+    });
     let compiledFn: CodegenFunction;
     try {
       compiledFn = compileFn(
@@ -430,11 +445,52 @@ export function compileProgram(
         prunedMemoValues: compiledFn.prunedMemoValues,
       });
     } catch (err) {
+      if (body.isBlockStatement()) {
+        const optOutDirectives = findDirectiveDisablingMemoization(body);
+        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 (body.isBlockStatement()) {
+      const optInDirectives = findDirectiveEnablingMemoization(body);
+      const optOutDirectives = findDirectiveDisablingMemoization(body);
+
+      /**
+       * Always compile functions with opt in directives.
+       */
+      if (optInDirectives.length > 0) {
+        return compiledFn;
+      } else if (pass.opts.compilationMode === 'annotation') {
+        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.node.value.value}' directive.`,
+            loc: directive.node.loc ?? null,
+          });
+        }
+        return null;
+      }
+    }
+
     if (!pass.opts.noEmit && !hasCriticalError) {
       return compiledFn;
     }
@@ -481,6 +537,12 @@ export function compileProgram(
     });
   }
 
+  const moduleScopeOptOutDirectives =
+    findDirectiveDisablingMemoization(program);
+  if (moduleScopeOptOutDirectives.length > 0) {
+    return;
+  }
+
   if (pass.opts.gating != null) {
     const error = checkFunctionReferencedBeforeDeclarationAtTopLevel(
       program,
@@ -596,24 +658,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;
@@ -630,29 +674,10 @@ function getReactFunctionType(
   environment: EnvironmentConfig,
 ): 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) {
+  const body = fn.get('body');
+  if (Array.isArray(body) === false && body.isBlockStatement()) {
+    if (findDirectiveEnablingMemoization(body).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/use-no-forget-multiple-with-eslint-suppression.js renamed to compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.use-no-forget-multiple-with-eslint-suppression.js

@@ -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 | }
+```
+          
+