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 15796d77a70e..48be507326f4 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 @@ -91,6 +98,7 @@ export class ProgramContext { filename, code, hasModuleScopeOptOut, + hasModuleScopeOptIn, }: ProgramContextOptions) { this.scope = program.scope; this.opts = opts; @@ -99,6 +107,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 2880e9283c77..b0c580ec5108 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts @@ -411,6 +411,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( @@ -507,7 +512,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; @@ -660,11 +669,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 { @@ -818,6 +829,7 @@ function shouldSkipCompilation( function getReactFunctionType( fn: BabelFn, pass: CompilerPass, + hasModuleScopeOptIn: boolean, ): ReactFunctionType | null { if (fn.node.body.type === 'BlockStatement') { const optInDirectives = tryFindDirectiveEnablingMemoization( @@ -841,7 +853,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}], +};