[compiler][rfc] Enable hook guards in dev mode by default

compiler/packages/babel-plugin-react-compiler/package.json

@@ -50,6 +50,7 @@
     "pretty-format": "^24",
     "react": "0.0.0-experimental-4beb1fd8-20241118",
     "react-dom": "0.0.0-experimental-4beb1fd8-20241118",
+    "react-compiler-runtime": "0.0.1",
     "ts-jest": "^29.1.1",
     "ts-node": "^10.9.2",
     "zod": "^3.22.4",

compiler/packages/babel-plugin-react-compiler/scripts/jest/makeTransform.ts

@@ -29,6 +29,7 @@ const e2eTransformerCacheKey = 1;
 const forgetOptions: EnvironmentConfig = validateEnvironmentConfig({
   enableAssumeHooksFollowRulesOfReact: true,
   enableFunctionOutlining: false,
+  enableEmitHookGuards: null,
 });
 const debugMode = process.env['DEBUG_FORGET_COMPILER'] != null;
 

compiler/packages/babel-plugin-react-compiler/src/Babel/BabelPlugin.ts

@@ -61,6 +61,18 @@ export default function BabelPluginReactCompiler(
               },
             };
           }
+          if (opts.environment.enableEmitHookGuards != null) {
+            const enableEmitHookGuards = opts.environment.enableEmitHookGuards;
+            if (enableEmitHookGuards.devonly === true && !isDev) {
+              opts = {
+                ...opts,
+                environment: {
+                  ...opts.environment,
+                  enableEmitHookGuards: null,
+                },
+              };
+            }
+          }
           compileProgram(prog, {
             opts,
             filename: pass.filename ?? null,

compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts

@@ -396,7 +396,27 @@ const EnvironmentConfigSchema = z.object({
    */
   enableEmitFreeze: ExternalFunctionSchema.nullable().default(null),
 
-  enableEmitHookGuards: ExternalFunctionSchema.nullable().default(null),
+  enableEmitHookGuards: z
+    .intersection(
+      ExternalFunctionSchema,
+      z.object({
+        devonly: z.union([
+          z.literal(true),
+          z.literal(false),
+          /**
+           * Backwards compatibility with previous configuration, which did not
+           * have this field.
+           */
+          z.literal(undefined),
+        ]),
+      }),
+    )
+    .nullable()
+    .default({
+      source: 'react-compiler-runtime',
+      importSpecifierName: '$dispatcherGuard',
+      devonly: true,
+    }),
 
   /**
    * Enable instruction reordering. See InstructionReordering.ts for the details
@@ -668,6 +688,7 @@ const testComplexConfigDefaults: PartialEnvironmentConfig = {
   enableEmitHookGuards: {
     source: 'react-compiler-runtime',
     importSpecifierName: '$dispatcherGuard',
+    devonly: false,
   },
   inlineJsxTransform: {
     elementSymbol: 'react.transitional.element',
@@ -711,6 +732,7 @@ function parseConfigPragmaEnvironmentForTest(
   const maybeConfig: any = {};
   // Get the defaults to programmatically check for boolean properties
   const defaultConfig = EnvironmentConfigSchema.parse({});
+  (maybeConfig as PartialEnvironmentConfig).enableEmitHookGuards = null;
 
   for (const token of pragma.split(' ')) {
     if (!token.startsWith('@')) {

compiler/packages/babel-plugin-react-compiler/src/__tests__/parseConfigPragma-test.ts

@@ -27,6 +27,7 @@ describe('parseConfigPragmaForTests()', () => {
       panicThreshold: 'all_errors',
       environment: {
         ...defaultOptions.environment,
+        enableEmitHookGuards: null,
         enableUseTypeAnnotations: true,
         validateNoSetStateInPassiveEffects: true,
         validateNoSetStateInRender: false,

compiler/packages/react-compiler-runtime/src/index.ts

@@ -41,7 +41,6 @@ export const c =
 
 const LazyGuardDispatcher: {[key: string]: (...args: Array<any>) => any} = {};
 [
-  'readContext',
   'useCallback',
   'useContext',
   'useEffect',
@@ -58,10 +57,8 @@ const LazyGuardDispatcher: {[key: string]: (...args: Array<any>) => any} = {};
   'useMutableSource',
   'useSyncExternalStore',
   'useId',
-  'unstable_isNewReconciler',
-  'getCacheSignal',
-  'getCacheForType',
   'useCacheRefresh',
+  'useOptimistic',
 ].forEach(name => {
   LazyGuardDispatcher[name] = () => {
     throw new Error(
@@ -74,15 +71,29 @@ const LazyGuardDispatcher: {[key: string]: (...args: Array<any>) => any} = {};
 let originalDispatcher: unknown = null;
 
 // Allow guards are not emitted for useMemoCache
-LazyGuardDispatcher['useMemoCache'] = (count: number) => {
-  if (originalDispatcher == null) {
-    throw new Error(
-      'React Compiler internal invariant violation: unexpected null dispatcher',
-    );
-  } else {
-    return (originalDispatcher as any).useMemoCache(count);
-  }
-};
+
+for (const key of [
+  // Allow guards may not be emitted for useMemoCache
+  'useMemoCache',
+  'unstable_useMemoCache',
+  // Not named as hooks
+  'readContext',
+  'unstable_isNewReconciler',
+  'getCacheSignal',
+  'getCacheForType',
+  // Not 'real' hooks (i.e. not implemented with an index based id)
+  'use',
+]) {
+  LazyGuardDispatcher[key] = (...args) => {
+    if (originalDispatcher == null) {
+      throw new Error(
+        'React Compiler internal invariant violation: unexpected null dispatcher',
+      );
+    } else {
+      return (originalDispatcher as any).useMemoCache(...args);
+    }
+  };
+}
 
 enum GuardKind {
   PushGuardContext = 0,
@@ -92,8 +103,13 @@ enum GuardKind {
 }
 
 function setCurrent(newDispatcher: any) {
-  ReactSecretInternals.ReactCurrentDispatcher.current = newDispatcher;
-  return ReactSecretInternals.ReactCurrentDispatcher.current;
+  if (ReactSecretInternals.ReactCurrentDispatcher != null) {
+    ReactSecretInternals.ReactCurrentDispatcher.current = newDispatcher;
+    return ReactSecretInternals.ReactCurrentDispatcher.current;
+  } else {
+    ReactSecretInternals.H = newDispatcher;
+    return ReactSecretInternals.H;
+  }
 }
 
 const guardFrames: Array<unknown> = [];
@@ -134,7 +150,9 @@ const guardFrames: Array<unknown> = [];
  * ```
  */
 export function $dispatcherGuard(kind: GuardKind) {
-  const curr = ReactSecretInternals.ReactCurrentDispatcher.current;
+  const curr =
+    ReactSecretInternals.H ??
+    ReactSecretInternals.ReactCurrentDispatcher.current;
   if (kind === GuardKind.PushGuardContext) {
     // Push before checking invariant or errors
     guardFrames.push(curr);
@@ -169,7 +187,29 @@ export function $dispatcherGuard(kind: GuardKind) {
     // ExpectHooks could be nested, so we save the current dispatcher
     // for the matching PopExpectHook to restore.
     guardFrames.push(curr);
-    setCurrent(originalDispatcher);
+    if (originalDispatcher != null) {
+      /**
+       * originalDispatcher could be null in the following case.
+       * ```js
+       * function Component() {
+       *   "use no memo";
+       *   const useFn = useMakeHook();
+       *   useFn();
+       * }
+       * function useMakeHook() {
+       *   // ...
+       *   return () => {
+       *     useState(...);
+       *   };
+       * }
+       * ```
+       *
+       * React compiler currently does not memoize within inner functions. While this
+       * code breaks the programming model (hooks as first class values), let's
+       * either ignore this or error with a better message (within another dispatcher)
+       */
+      setCurrent(originalDispatcher);
+    }
   } else if (kind === GuardKind.PopExpectHook) {
     const lastFrame = guardFrames.pop();
     if (lastFrame == null) {