[compiler] Fix module-level 'use memo' not compiling in annotation mode

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

@@ -60,6 +60,7 @@ type ProgramContextOptions = {
   filename: string | null;
   code: string | null;
   hasModuleScopeOptOut: boolean;
+  hasModuleScopeOptIn: boolean;
 };
 export class ProgramContext {
   /**
@@ -72,6 +73,12 @@ export class ProgramContext {
   reactRuntimeModule: string;
   suppressions: Array<SuppressionRange>;
   hasModuleScopeOptOut: boolean;
+  /**
+   * True when a module-level opt-in directive (e.g. `'use memo'`) is present.
+   * In annotation mode this makes every component/hook in the module behave as
+   * if it had an individual function-level opt-in directive.
+   */
+  hasModuleScopeOptIn: boolean;
 
   /*
    * This is a hack to work around what seems to be a Babel bug. Babel doesn't
@@ -96,6 +103,7 @@ export class ProgramContext {
     filename,
     code,
     hasModuleScopeOptOut,
+    hasModuleScopeOptIn,
   }: ProgramContextOptions) {
     this.scope = program.scope;
     this.opts = opts;
@@ -104,6 +112,7 @@ export class ProgramContext {
     this.reactRuntimeModule = getReactCompilerRuntimeModule(opts.target);
     this.suppressions = suppressions;
     this.hasModuleScopeOptOut = hasModuleScopeOptOut;
+    this.hasModuleScopeOptIn = hasModuleScopeOptIn;
   }
 
   isHookName(name: string): boolean {

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

@@ -414,6 +414,11 @@ export function compileProgram(
     hasModuleScopeOptOut:
       findDirectiveDisablingMemoization(program.node.directives, pass.opts) !=
       null,
+    hasModuleScopeOptIn:
+      tryFindDirectiveEnablingMemoization(
+        program.node.directives,
+        pass.opts,
+      ).unwrapOr(null) != null,
   });
 
   const queue: Array<CompileSource> = findFunctionsToCompile(
@@ -514,7 +519,11 @@ function findFunctionsToCompile(
       return;
     }
 
-    const fnType = getReactFunctionType(fn, pass);
+    const fnType = getReactFunctionType(
+      fn,
+      pass,
+      programContext.hasModuleScopeOptIn,
+    );
 
     if (fnType === null || programContext.alreadyCompiled.has(fn.node)) {
       return;
@@ -667,11 +676,13 @@ function processFn(
     return null;
   } else if (
     programContext.opts.compilationMode === 'annotation' &&
-    directives.optIn == null
+    directives.optIn == null &&
+    !programContext.hasModuleScopeOptIn
   ) {
     /**
-     * If no opt-in directive is found and the compiler is configured in
-     * annotation mode, don't insert the compiled function.
+     * If no opt-in directive is found (neither function-level nor module-level)
+     * and the compiler is configured in annotation mode, don't insert the
+     * compiled function.
      */
     return null;
   } else {
@@ -809,6 +820,7 @@ function shouldSkipCompilation(
 function getReactFunctionType(
   fn: BabelFn,
   pass: CompilerPass,
+  hasModuleScopeOptIn: boolean,
 ): ReactFunctionType | null {
   if (fn.node.body.type === 'BlockStatement') {
     const optInDirectives = tryFindDirectiveEnablingMemoization(
@@ -832,7 +844,11 @@ function getReactFunctionType(
 
   switch (pass.opts.compilationMode) {
     case 'annotation': {
-      // opt-ins are checked above
+      // opt-ins are checked above (function-level)
+      // A module-level opt-in directive applies to all functions in the file
+      if (hasModuleScopeOptIn) {
+        return getComponentOrHookLike(fn) ?? 'Other';
+      }
       return null;
     }
     case 'infer': {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/use-memo-module-level.expect.md

@@ -0,0 +1,49 @@
+
+## Input
+
+```javascript
+// @compilationMode:"annotation"
+'use memo';
+
+function Component({a, b}) {
+  return <div>{a + b}</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 1, b: 2}],
+};
+
+```
+
+## Code
+
+```javascript
+// @compilationMode:"annotation"
+"use memo";
+import { c as _c } from "react/compiler-runtime";
+
+function Component(t0) {
+  const $ = _c(2);
+  const { a, b } = t0;
+  const t1 = a + b;
+  let t2;
+  if ($[0] !== t1) {
+    t2 = <div>{t1}</div>;
+    $[0] = t1;
+    $[1] = t2;
+  } else {
+    t2 = $[1];
+  }
+  return t2;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ a: 1, b: 2 }],
+};
+
+```
+
+### Eval output
+(kind: ok) <div>3</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/use-memo-module-level.js

@@ -0,0 +1,11 @@
+// @compilationMode:"annotation"
+'use memo';
+
+function Component({a, b}) {
+  return <div>{a + b}</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 1, b: 2}],
+};