From 4a3886eb3a1e56bda51d4908c0e9c48f103c222a Mon Sep 17 00:00:00 2001 From: Dmitrii Troitskii Date: Mon, 23 Feb 2026 13:13:00 +0000 Subject: [PATCH] [compiler] Fix module-level 'use memo' not compiling in annotation mode When compilationMode is set to 'annotation', module-level opt-in directives such as 'use memo' (placed at the top of a file) were silently ignored. Only function-level directives triggered compilation, contrary to the documented behavior that "place directives at the top of a file to affect all functions in that module". Root cause: both getReactFunctionType() and processFn() only inspected the function body's own directives (fn.node.body.directives) when deciding whether to compile in annotation mode. The program-level directives (program.node.directives) were never consulted. Fix: - Add hasModuleScopeOptIn to ProgramContext, computed once from tryFindDirectiveEnablingMemoization(program.node.directives). - In getReactFunctionType(), when compilationMode is 'annotation' and no function-level opt-in is found, fall back to hasModuleScopeOptIn before returning null. This ensures functions in the module are queued for compilation. - In processFn(), add !programContext.hasModuleScopeOptIn to the annotation mode skip guard so that the compiled function is emitted. Fixes facebook/react#35868 --- .../src/Entrypoint/Imports.ts | 9 ++++ .../src/Entrypoint/Program.ts | 26 ++++++++-- .../compiler/use-memo-module-level.expect.md | 49 +++++++++++++++++++ .../compiler/use-memo-module-level.js | 11 +++++ 4 files changed, 90 insertions(+), 5 deletions(-) create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/use-memo-module-level.expect.md create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/use-memo-module-level.js diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Imports.ts b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Imports.ts index 2fef4cfabe59..73316d672e64 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Imports.ts +++ b/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; 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 { diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts index de36ad218f7e..556fc9a517b2 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts +++ b/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 = 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': { diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/use-memo-module-level.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/use-memo-module-level.expect.md new file mode 100644 index 000000000000..22ba70b215b0 --- /dev/null +++ b/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
{a + b}
; +} + +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 =
{t1}
; + $[0] = t1; + $[1] = t2; + } else { + t2 = $[1]; + } + return t2; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ a: 1, b: 2 }], +}; + +``` + +### Eval output +(kind: ok)
3
\ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/use-memo-module-level.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/use-memo-module-level.js new file mode 100644 index 000000000000..43d80e3223d9 --- /dev/null +++ b/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
{a + b}
; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{a: 1, b: 2}], +};