[compiler][ez] Clean up pragma parsing for tests + playground (#31347)

Move environment config parsing for `inlineJsxTransform`,
`lowerContextAccess`, and some dev-only options out of snap (test
fixture). These should now be available for playground via
`@inlineJsxTransform` and `lowerContextAccess`.

Other small change:
Changed zod fields from `nullish()` -> `nullable().default(null)`.
[`nullish`](https://zod.dev/?id=nullish) fields accept `null |
undefined` and default to `undefined`. We don't distinguish between null
and undefined for any of these options, so let's only accept null +
default to null. This also makes EnvironmentConfig in the playground
more accurate. Previously, some fields just didn't show up as
`prettyFormat({field: undefined})` does not print `field`.

compiler/apps/playground/components/Editor/EditorImpl.tsx

@@ -14,7 +14,7 @@ import {
   CompilerErrorDetail,
   Effect,
   ErrorSeverity,
-  parseConfigPragma,
+  parseConfigPragmaForTests,
   ValueKind,
   runPlayground,
   type Hook,
@@ -208,7 +208,7 @@ function compile(source: string): [CompilerOutput, 'flow' | 'typescript'] {
   try {
     // Extract the first line to quickly check for custom test directives
     const pragma = source.substring(0, source.indexOf('\n'));
-    const config = parseConfigPragma(pragma);
+    const config = parseConfigPragmaForTests(pragma);
 
     for (const fn of parseFunctions(source, language)) {
       const id = withIdentifier(getFunctionIdentifier(fn));

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

@@ -69,8 +69,8 @@ export const ExternalFunctionSchema = z.object({
 export const InstrumentationSchema = z
   .object({
     fn: ExternalFunctionSchema,
-    gating: ExternalFunctionSchema.nullish(),
-    globalGating: z.string().nullish(),
+    gating: ExternalFunctionSchema.nullable(),
+    globalGating: z.string().nullable(),
   })
   .refine(
     opts => opts.gating != null || opts.globalGating != null,
@@ -147,7 +147,7 @@ export type Hook = z.infer<typeof HookSchema>;
  */
 
 const EnvironmentConfigSchema = z.object({
-  customHooks: z.map(z.string(), HookSchema).optional().default(new Map()),
+  customHooks: z.map(z.string(), HookSchema).default(new Map()),
 
   /**
    * A function that, given the name of a module, can optionally return a description
@@ -248,7 +248,7 @@ const EnvironmentConfigSchema = z.object({
    *
    * The symbol configuration is set for backwards compatability with pre-React 19 transforms
    */
-  inlineJsxTransform: ReactElementSymbolSchema.nullish(),
+  inlineJsxTransform: ReactElementSymbolSchema.nullable().default(null),
 
   /*
    * Enable validation of hooks to partially check that the component honors the rules of hooks.
@@ -339,9 +339,9 @@ const EnvironmentConfigSchema = z.object({
    *     }
    *   }
    */
-  enableEmitFreeze: ExternalFunctionSchema.nullish(),
+  enableEmitFreeze: ExternalFunctionSchema.nullable().default(null),
 
-  enableEmitHookGuards: ExternalFunctionSchema.nullish(),
+  enableEmitHookGuards: ExternalFunctionSchema.nullable().default(null),
 
   /**
    * Enable instruction reordering. See InstructionReordering.ts for the details
@@ -425,7 +425,7 @@ const EnvironmentConfigSchema = z.object({
    *   }
    *
    */
-  enableEmitInstrumentForget: InstrumentationSchema.nullish(),
+  enableEmitInstrumentForget: InstrumentationSchema.nullable().default(null),
 
   // Enable validation of mutable ranges
   assertValidMutableRanges: z.boolean().default(false),
@@ -464,8 +464,6 @@ const EnvironmentConfigSchema = z.object({
    */
   throwUnknownException__testonly: z.boolean().default(false),
 
-  enableSharedRuntime__testonly: z.boolean().default(false),
-
   /**
    * Enables deps of a function epxression to be treated as conditional. This
    * makes sure we don't load a dep when it's a property (to check if it has
@@ -503,7 +501,8 @@ const EnvironmentConfigSchema = z.object({
    * computed one. This detects cases where rules of react violations may cause the
    * compiled code to behave differently than the original.
    */
-  enableChangeDetectionForDebugging: ExternalFunctionSchema.nullish(),
+  enableChangeDetectionForDebugging:
+    ExternalFunctionSchema.nullable().default(null),
 
   /**
    * The react native re-animated library uses custom Babel transforms that
@@ -543,7 +542,7 @@ const EnvironmentConfigSchema = z.object({
    *
    * Here the variables `ref` and `myRef` will be typed as Refs.
    */
-  enableTreatRefLikeIdentifiersAsRefs: z.boolean().nullable().default(false),
+  enableTreatRefLikeIdentifiersAsRefs: z.boolean().default(false),
 
   /*
    * If specified a value, the compiler lowers any calls to `useContext` to use
@@ -565,12 +564,57 @@ const EnvironmentConfigSchema = z.object({
    * const {foo, bar} = useCompiledContext(MyContext, (c) => [c.foo, c.bar]);
    * ```
    */
-  lowerContextAccess: ExternalFunctionSchema.nullish(),
+  lowerContextAccess: ExternalFunctionSchema.nullable().default(null),
 });
 
 export type EnvironmentConfig = z.infer<typeof EnvironmentConfigSchema>;
 
-export function parseConfigPragma(pragma: string): EnvironmentConfig {
+/**
+ * For test fixtures and playground only.
+ *
+ * Pragmas are straightforward to parse for boolean options (`:true` and
+ * `:false`). These are 'enabled' config values for non-boolean configs (i.e.
+ * what is used when parsing `:true`).
+ */
+const testComplexConfigDefaults: PartialEnvironmentConfig = {
+  validateNoCapitalizedCalls: [],
+  enableChangeDetectionForDebugging: {
+    source: 'react-compiler-runtime',
+    importSpecifierName: '$structuralCheck',
+  },
+  enableEmitFreeze: {
+    source: 'react-compiler-runtime',
+    importSpecifierName: 'makeReadOnly',
+  },
+  enableEmitInstrumentForget: {
+    fn: {
+      source: 'react-compiler-runtime',
+      importSpecifierName: 'useRenderCounter',
+    },
+    gating: {
+      source: 'react-compiler-runtime',
+      importSpecifierName: 'shouldInstrument',
+    },
+    globalGating: '__DEV__',
+  },
+  enableEmitHookGuards: {
+    source: 'react-compiler-runtime',
+    importSpecifierName: '$dispatcherGuard',
+  },
+  inlineJsxTransform: {
+    elementSymbol: 'react.transitional.element',
+    globalDevVar: 'DEV',
+  },
+  lowerContextAccess: {
+    source: 'react-compiler-runtime',
+    importSpecifierName: 'useContext_withSelector',
+  },
+};
+
+/**
+ * For snap test fixtures and playground only.
+ */
+export function parseConfigPragmaForTests(pragma: string): EnvironmentConfig {
   const maybeConfig: any = {};
   // Get the defaults to programmatically check for boolean properties
   const defaultConfig = EnvironmentConfigSchema.parse({});
@@ -580,21 +624,12 @@ export function parseConfigPragma(pragma: string): EnvironmentConfig {
       continue;
     }
     const keyVal = token.slice(1);
-    let [key, val]: any = keyVal.split(':');
-
-    if (key === 'validateNoCapitalizedCalls') {
-      maybeConfig[key] = [];
-      continue;
-    }
+    let [key, val = undefined] = keyVal.split(':');
+    const isSet = val === undefined || val === 'true';
 
-    if (
-      key === 'enableChangeDetectionForDebugging' &&
-      (val === undefined || val === 'true')
-    ) {
-      maybeConfig[key] = {
-        source: 'react-compiler-runtime',
-        importSpecifierName: '$structuralCheck',
-      };
+    if (isSet && key in testComplexConfigDefaults) {
+      maybeConfig[key] =
+        testComplexConfigDefaults[key as keyof PartialEnvironmentConfig];
       continue;
     }
 
@@ -609,7 +644,6 @@ export function parseConfigPragma(pragma: string): EnvironmentConfig {
             props.push({type: 'name', name: elt});
           }
         }
-        console.log([valSplit[0], props.map(x => x.name ?? '*').join('.')]);
         maybeConfig[key] = [[valSplit[0], props]];
       }
       continue;
@@ -620,11 +654,10 @@ export function parseConfigPragma(pragma: string): EnvironmentConfig {
       continue;
     }
     if (val === undefined || val === 'true') {
-      val = true;
+      maybeConfig[key] = true;
     } else {
-      val = false;
+      maybeConfig[key] = false;
     }
-    maybeConfig[key] = val;
   }
 
   const config = EnvironmentConfigSchema.safeParse(maybeConfig);

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

@@ -17,7 +17,7 @@ export {buildReactiveScopeTerminalsHIR} from './BuildReactiveScopeTerminalsHIR';
 export {computeDominatorTree, computePostDominatorTree} from './Dominator';
 export {
   Environment,
-  parseConfigPragma,
+  parseConfigPragmaForTests,
   validateEnvironmentConfig,
   type EnvironmentConfig,
   type ExternalFunction,

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-emit-imports-same-source.expect.md

@@ -2,7 +2,7 @@
 ## Input
 
 ```javascript
-// @enableEmitFreeze @instrumentForget
+// @enableEmitFreeze @enableEmitInstrumentForget
 
 function useFoo(props) {
   return foo(props.x);
@@ -18,7 +18,7 @@ import {
   shouldInstrument,
   makeReadOnly,
 } from "react-compiler-runtime";
-import { c as _c } from "react/compiler-runtime"; // @enableEmitFreeze @instrumentForget
+import { c as _c } from "react/compiler-runtime"; // @enableEmitFreeze @enableEmitInstrumentForget
 
 function useFoo(props) {
   if (__DEV__ && shouldInstrument)

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-emit-imports-same-source.js

@@ -1,4 +1,4 @@
-// @enableEmitFreeze @instrumentForget
+// @enableEmitFreeze @enableEmitInstrumentForget
 
 function useFoo(props) {
   return foo(props.x);

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-instrument-forget-gating-test.expect.md

@@ -2,7 +2,7 @@
 ## Input
 
 ```javascript
-// @instrumentForget @compilationMode(annotation) @gating
+// @enableEmitInstrumentForget @compilationMode(annotation) @gating
 
 function Bar(props) {
   'use forget';
@@ -25,7 +25,7 @@ function Foo(props) {
 ```javascript
 import { isForgetEnabled_Fixtures } from "ReactForgetFeatureFlag";
 import { useRenderCounter, shouldInstrument } from "react-compiler-runtime";
-import { c as _c } from "react/compiler-runtime"; // @instrumentForget @compilationMode(annotation) @gating
+import { c as _c } from "react/compiler-runtime"; // @enableEmitInstrumentForget @compilationMode(annotation) @gating
 const Bar = isForgetEnabled_Fixtures()
   ? function Bar(props) {
       "use forget";

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-instrument-forget-gating-test.js

@@ -1,4 +1,4 @@
-// @instrumentForget @compilationMode(annotation) @gating
+// @enableEmitInstrumentForget @compilationMode(annotation) @gating
 
 function Bar(props) {
   'use forget';

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-instrument-forget-test.expect.md

@@ -2,7 +2,7 @@
 ## Input
 
 ```javascript
-// @instrumentForget @compilationMode(annotation)
+// @enableEmitInstrumentForget @compilationMode(annotation)
 
 function Bar(props) {
   'use forget';
@@ -24,7 +24,7 @@ function Foo(props) {
 
 ```javascript
 import { useRenderCounter, shouldInstrument } from "react-compiler-runtime";
-import { c as _c } from "react/compiler-runtime"; // @instrumentForget @compilationMode(annotation)
+import { c as _c } from "react/compiler-runtime"; // @enableEmitInstrumentForget @compilationMode(annotation)
 
 function Bar(props) {
   "use forget";

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-instrument-forget-test.js

@@ -1,4 +1,4 @@
-// @instrumentForget @compilationMode(annotation)
+// @enableEmitInstrumentForget @compilationMode(annotation)
 
 function Bar(props) {
   'use forget';

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.expect.md

@@ -2,7 +2,7 @@
 ## Input
 
 ```javascript
-// @enableInlineJsxTransform
+// @inlineJsxTransform
 
 function Parent({children, a: _a, b: _b, c: _c, ref}) {
   return <div ref={ref}>{children}</div>;
@@ -76,7 +76,7 @@ export const FIXTURE_ENTRYPOINT = {
 ## Code
 
 ```javascript
-import { c as _c2 } from "react/compiler-runtime"; // @enableInlineJsxTransform
+import { c as _c2 } from "react/compiler-runtime"; // @inlineJsxTransform
 
 function Parent(t0) {
   const $ = _c2(2);

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.js

@@ -1,4 +1,4 @@
-// @enableInlineJsxTransform
+// @inlineJsxTransform
 
 function Parent({children, a: _a, b: _b, c: _c, ref}) {
   return <div ref={ref}>{children}</div>;

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

@@ -5,9 +5,9 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import {parseConfigPragma, validateEnvironmentConfig} from '..';
+import {parseConfigPragmaForTests, validateEnvironmentConfig} from '..';
 
-describe('parseConfigPragma()', () => {
+describe('parseConfigPragmaForTests()', () => {
   it('parses flags in various forms', () => {
     const defaultConfig = validateEnvironmentConfig({});
 
@@ -17,7 +17,7 @@ describe('parseConfigPragma()', () => {
     expect(defaultConfig.validateNoSetStateInPassiveEffects).toBe(false);
     expect(defaultConfig.validateNoSetStateInRender).toBe(true);
 
-    const config = parseConfigPragma(
+    const config = parseConfigPragmaForTests(
       '@enableUseTypeAnnotations @validateNoSetStateInPassiveEffects:true @validateNoSetStateInRender:false',
     );
     expect(config).toEqual({

compiler/packages/babel-plugin-react-compiler/src/index.ts

@@ -26,7 +26,7 @@ export {
 export {
   Effect,
   ValueKind,
-  parseConfigPragma,
+  parseConfigPragmaForTests,
   printHIR,
   validateEnvironmentConfig,
   type EnvironmentConfig,