[compiler][gating] Experimental directive based gating

Adds `dynamicGating` as an experimental option for testing rollout DX at Meta. If specified, this enables dynamic gating which matches `use memo if(...)` directives.

#### Example usage
Input file
```js
// @dynamicGating:{"source":"myModule"}
export function MyComponent() {
  'use memo if(isEnabled)';
   return <div>...</div>;
}
```
Compiler output
```js
import {isEnabled} from 'myModule';
export const MyComponent = isEnabled()
  ? <optimized version>
  : <original version>;
```

Author: mofeiZ

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

@@ -37,6 +37,10 @@ const PanicThresholdOptionsSchema = z.enum([
 ]);
 
 export type PanicThresholdOptions = z.infer<typeof PanicThresholdOptionsSchema>;
+const DynamicGatingOptionsSchema = z.object({
+  source: z.string(),
+});
+export type DynamicGatingOptions = z.infer<typeof DynamicGatingOptionsSchema>;
 
 export type PluginOptions = {
   environment: EnvironmentConfig;
@@ -65,6 +69,28 @@ export type PluginOptions = {
    */
   gating: ExternalFunction | null;
 
+  /**
+   * If specified, this enables dynamic gating which matches `use memo if(...)`
+   * directives.
+   *
+   * Example usage:
+   * ```js
+   * // @dynamicGating:{"source":"myModule"}
+   * export function MyComponent() {
+   *   'use memo if(isEnabled)';
+   *    return <div>...</div>;
+   * }
+   * ```
+   * This will emit:
+   * ```js
+   * import {isEnabled} from 'myModule';
+   * export const MyComponent = isEnabled()
+   *   ? <optimized version>
+   *   : <original version>;
+   * ```
+   */
+  dynamicGating: DynamicGatingOptions | null;
+
   panicThreshold: PanicThresholdOptions;
 
   /*
@@ -244,6 +270,7 @@ export const defaultOptions: PluginOptions = {
   logger: null,
   gating: null,
   noEmit: false,
+  dynamicGating: null,
   eslintSuppressionRules: null,
   flowSuppressions: true,
   ignoreUseNoForget: false,
@@ -292,6 +319,25 @@ export function parsePluginOptions(obj: unknown): PluginOptions {
           }
           break;
         }
+        case 'dynamicGating': {
+          if (value == null) {
+            parsedOptions[key] = null;
+          } else {
+            const result = DynamicGatingOptionsSchema.safeParse(value);
+            if (result.success) {
+              parsedOptions[key] = result.data;
+            } else {
+              CompilerError.throwInvalidConfig({
+                reason:
+                  'Could not parse dynamic gating. Update React Compiler config to fix the error',
+                description: `${fromZodError(result.error)}`,
+                loc: null,
+                suggestions: null,
+              });
+            }
+          }
+          break;
+        }
         default: {
           parsedOptions[key] = value;
         }

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

@@ -12,7 +12,7 @@ import {
   CompilerErrorDetail,
   ErrorSeverity,
 } from '../CompilerError';
-import {ReactFunctionType} from '../HIR/Environment';
+import {ExternalFunction, ReactFunctionType} from '../HIR/Environment';
 import {CodegenFunction} from '../ReactiveScopes';
 import {isComponentDeclaration} from '../Utils/ComponentDeclaration';
 import {isHookDeclaration} from '../Utils/HookDeclaration';
@@ -31,6 +31,7 @@ import {
   suppressionsToCompilerError,
 } from './Suppression';
 import {GeneratedSource} from '../HIR';
+import {Err, Ok, Result} from '../Utils/Result';
 
 export type CompilerPass = {
   opts: PluginOptions;
@@ -40,15 +41,24 @@ export type CompilerPass = {
 };
 export const OPT_IN_DIRECTIVES = new Set(['use forget', 'use memo']);
 export const OPT_OUT_DIRECTIVES = new Set(['use no forget', 'use no memo']);
+const DYNAMIC_GATING_DIRECTIVE = new RegExp('^use memo if\\(([^\\)]*)\\)$');
 
-export function findDirectiveEnablingMemoization(
+export function tryFindDirectiveEnablingMemoization(
   directives: Array<t.Directive>,
-): t.Directive | null {
-  return (
-    directives.find(directive =>
-      OPT_IN_DIRECTIVES.has(directive.value.value),
-    ) ?? null
+  opts: PluginOptions,
+): Result<t.Directive | null, CompilerError> {
+  const optIn = directives.find(directive =>
+    OPT_IN_DIRECTIVES.has(directive.value.value),
   );
+  if (optIn != null) {
+    return Ok(optIn);
+  }
+  const dynamicGating = findDirectivesDynamicGating(directives, opts);
+  if (dynamicGating.isOk()) {
+    return Ok(dynamicGating.unwrap()?.directive ?? null);
+  } else {
+    return Err(dynamicGating.unwrapErr());
+  }
 }
 
 export function findDirectiveDisablingMemoization(
@@ -60,6 +70,64 @@ export function findDirectiveDisablingMemoization(
     ) ?? null
   );
 }
+function findDirectivesDynamicGating(
+  directives: Array<t.Directive>,
+  opts: PluginOptions,
+): Result<
+  {
+    gating: ExternalFunction;
+    directive: t.Directive;
+  } | null,
+  CompilerError
+> {
+  if (opts.dynamicGating === null) {
+    return Ok(null);
+  }
+  const errors = new CompilerError();
+  const result: Array<{directive: t.Directive; match: string}> = [];
+
+  for (const directive of directives) {
+    const maybeMatch = DYNAMIC_GATING_DIRECTIVE.exec(directive.value.value);
+    if (maybeMatch != null && maybeMatch[1] != null) {
+      if (t.isValidIdentifier(maybeMatch[1])) {
+        result.push({directive, match: maybeMatch[1]});
+      } else {
+        errors.push({
+          reason: `Dynamic gating directive is not a valid JavaScript identifier`,
+          description: `Found '${directive.value.value}'`,
+          severity: ErrorSeverity.InvalidReact,
+          loc: directive.loc ?? null,
+          suggestions: null,
+        });
+      }
+    }
+  }
+  if (errors.hasErrors()) {
+    return Err(errors);
+  } else if (result.length > 1) {
+    const error = new CompilerError();
+    error.push({
+      reason: `Multiple dynamic gating directives found`,
+      description: `Expected a single directive but found [${result
+        .map(r => r.directive.value.value)
+        .join(', ')}]`,
+      severity: ErrorSeverity.InvalidReact,
+      loc: result[0].directive.loc ?? null,
+      suggestions: null,
+    });
+    return Err(error);
+  } else if (result.length === 1) {
+    return Ok({
+      gating: {
+        source: opts.dynamicGating.source,
+        importSpecifierName: result[0].match,
+      },
+      directive: result[0].directive,
+    });
+  } else {
+    return Ok(null);
+  }
+}
 
 function isCriticalError(err: unknown): boolean {
   return !(err instanceof CompilerError) || err.isCritical();
@@ -477,12 +545,32 @@ function processFn(
   fnType: ReactFunctionType,
   programContext: ProgramContext,
 ): null | CodegenFunction {
-  let directives;
+  let directives: {
+    optIn: t.Directive | null;
+    optOut: t.Directive | null;
+  };
   if (fn.node.body.type !== 'BlockStatement') {
-    directives = {optIn: null, optOut: null};
+    directives = {
+      optIn: null,
+      optOut: null,
+    };
   } else {
+    const optIn = tryFindDirectiveEnablingMemoization(
+      fn.node.body.directives,
+      programContext.opts,
+    );
+    if (optIn.isErr()) {
+      /**
+       * If parsing opt-in directive fails, it's most likely that React Compiler
+       * was not tested or rolled out on this function. In that case, we handle
+       * the error and fall back to the safest option which is to not optimize
+       * the function.
+       */
+      handleError(optIn.unwrapErr(), programContext, fn.node.loc ?? null);
+      return null;
+    }
     directives = {
-      optIn: findDirectiveEnablingMemoization(fn.node.body.directives),
+      optIn: optIn.unwrapOr(null),
       optOut: findDirectiveDisablingMemoization(fn.node.body.directives),
     };
   }
@@ -659,25 +747,31 @@ function applyCompiledFunctions(
   pass: CompilerPass,
   programContext: ProgramContext,
 ): void {
-  const referencedBeforeDeclared =
-    pass.opts.gating != null
-      ? getFunctionReferencedBeforeDeclarationAtTopLevel(program, compiledFns)
-      : null;
+  let referencedBeforeDeclared = null;
   for (const result of compiledFns) {
     const {kind, originalFn, compiledFn} = result;
     const transformedFn = createNewFunctionNode(originalFn, compiledFn);
     programContext.alreadyCompiled.add(transformedFn);
 
-    if (referencedBeforeDeclared != null && kind === 'original') {
-      CompilerError.invariant(pass.opts.gating != null, {
-        reason: "Expected 'gating' import to be present",
-        loc: null,
-      });
+    let dynamicGating: ExternalFunction | null = null;
+    if (originalFn.node.body.type === 'BlockStatement') {
+      const result = findDirectivesDynamicGating(
+        originalFn.node.body.directives,
+        pass.opts,
+      );
+      if (result.isOk()) {
+        dynamicGating = result.unwrap()?.gating ?? null;
+      }
+    }
+    const functionGating = dynamicGating ?? pass.opts.gating;
+    if (kind === 'original' && functionGating != null) {
+      referencedBeforeDeclared ??=
+        getFunctionReferencedBeforeDeclarationAtTopLevel(program, compiledFns);
       insertGatedFunctionDeclaration(
         originalFn,
         transformedFn,
         programContext,
-        pass.opts.gating,
+        functionGating,
         referencedBeforeDeclared.has(result),
       );
     } else {
@@ -733,8 +827,13 @@ function getReactFunctionType(
 ): ReactFunctionType | null {
   const hookPattern = pass.opts.environment.hookPattern;
   if (fn.node.body.type === 'BlockStatement') {
-    if (findDirectiveEnablingMemoization(fn.node.body.directives) != null)
+    const optInDirectives = tryFindDirectiveEnablingMemoization(
+      fn.node.body.directives,
+      pass.opts,
+    );
+    if (optInDirectives.unwrapOr(null) != null) {
       return getComponentOrHookLike(fn, hookPattern) ?? 'Other';
+    }
   }
 
   // Component and hook declarations are known components/hooks

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/gating/dynamic-gating-annotation.expect.md

@@ -0,0 +1,50 @@
+
+## Input
+
+```javascript
+// @dynamicGating:{"source":"shared-runtime"} @compilationMode:"annotation"
+
+function Foo() {
+  'use memo if(getTrue)';
+  return <div>hello world</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { getTrue } from "shared-runtime"; // @dynamicGating:{"source":"shared-runtime"} @compilationMode:"annotation"
+const Foo = getTrue()
+  ? function Foo() {
+      "use memo if(getTrue)";
+      const $ = _c(1);
+      let t0;
+      if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+        t0 = <div>hello world</div>;
+        $[0] = t0;
+      } else {
+        t0 = $[0];
+      }
+      return t0;
+    }
+  : function Foo() {
+      "use memo if(getTrue)";
+      return <div>hello world</div>;
+    };
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{}],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>hello world</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/gating/dynamic-gating-annotation.js

@@ -0,0 +1,11 @@
+// @dynamicGating:{"source":"shared-runtime"} @compilationMode:"annotation"
+
+function Foo() {
+  'use memo if(getTrue)';
+  return <div>hello world</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{}],
+};