From e88b05e387a5793a33c9541f3f7bc749befd2c1f Mon Sep 17 00:00:00 2001
From: Mofei Zhang <feifei0@meta.com>
Date: Tue, 5 Nov 2024 15:15:40 -0500
Subject: [PATCH] [compiler][ez] Clean up pragma parsing for tests + playground

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`.
---
 .../src/HIR/Environment.ts                    | 93 +++++++++++++------
 ...codegen-emit-imports-same-source.expect.md |  4 +-
 .../codegen-emit-imports-same-source.js       |  2 +-
 ...en-instrument-forget-gating-test.expect.md |  4 +-
 .../codegen-instrument-forget-gating-test.js  |  2 +-
 .../codegen-instrument-forget-test.expect.md  |  4 +-
 .../codegen-instrument-forget-test.js         |  2 +-
 .../compiler/inline-jsx-transform.expect.md   |  4 +-
 .../fixtures/compiler/inline-jsx-transform.js |  2 +-
 compiler/packages/snap/src/compiler.ts        | 72 +++-----------
 10 files changed, 88 insertions(+), 101 deletions(-)

diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts
index 3e2b5597ac..8130e46631 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts
+++ b/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,11 +564,56 @@ 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>;
 
+/**
+ * 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 testConfigValues: 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 parseConfigPragma(pragma: string): EnvironmentConfig {
   const maybeConfig: any = {};
   // Get the defaults to programmatically check for boolean properties
@@ -580,21 +624,12 @@ export function parseConfigPragma(pragma: string): EnvironmentConfig {
       continue;
     }
     const keyVal = token.slice(1);
-    let [key, val]: any = keyVal.split(':');
+    let [key, val = undefined] = keyVal.split(':');
+    const isSet = val === undefined || val === 'true';
 
-    if (key === 'validateNoCapitalizedCalls') {
-      maybeConfig[key] = [];
-      continue;
-    }
-
-    if (
-      key === 'enableChangeDetectionForDebugging' &&
-      (val === undefined || val === 'true')
-    ) {
-      maybeConfig[key] = {
-        source: 'react-compiler-runtime',
-        importSpecifierName: '$structuralCheck',
-      };
+    if (isSet && key in testConfigValues) {
+      maybeConfig[key] =
+        testConfigValues[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);
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-emit-imports-same-source.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-emit-imports-same-source.expect.md
index dd67bcfbff..9a59b36cc0 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-emit-imports-same-source.expect.md
+++ b/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)
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-emit-imports-same-source.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-emit-imports-same-source.js
index 4edff1c3fc..bd66353319 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-emit-imports-same-source.js
+++ b/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);
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-instrument-forget-gating-test.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-instrument-forget-gating-test.expect.md
index 4aa29992eb..fc9247344d 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-instrument-forget-gating-test.expect.md
+++ b/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";
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-instrument-forget-gating-test.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-instrument-forget-gating-test.js
index 85fbd97ee7..dffb8ce795 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-instrument-forget-gating-test.js
+++ b/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';
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-instrument-forget-test.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-instrument-forget-test.expect.md
index ba8ed5056b..b5da853b6e 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-instrument-forget-test.expect.md
+++ b/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";
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-instrument-forget-test.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-instrument-forget-test.js
index 8947503277..2aef527e6b 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/codegen-instrument-forget-test.js
+++ b/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';
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.expect.md
index e657e36d36..f622b3a6fd 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.expect.md
+++ b/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);
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.js
index bebb7ad53b..ca55cab4ff 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.js
+++ b/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>;
diff --git a/compiler/packages/snap/src/compiler.ts b/compiler/packages/snap/src/compiler.ts
index f0ee88f06e..74737f96a0 100644
--- a/compiler/packages/snap/src/compiler.ts
+++ b/compiler/packages/snap/src/compiler.ts
@@ -21,7 +21,6 @@ import type {
 } from 'babel-plugin-react-compiler/src/Entrypoint';
 import type {Effect, ValueKind} from 'babel-plugin-react-compiler/src/HIR';
 import type {
-  EnvironmentConfig,
   Macro,
   MacroMethod,
   parseConfigPragma as ParseConfigPragma,
@@ -37,6 +36,11 @@ export function parseLanguage(source: string): 'flow' | 'typescript' {
   return source.indexOf('@flow') !== -1 ? 'flow' : 'typescript';
 }
 
+/**
+ * Parse react compiler plugin + environment options from test fixture. Note
+ * that although this primarily uses `Environment:parseConfigPragma`, it also
+ * has test fixture specific (i.e. not applicable to playground) parsing logic.
+ */
 function makePluginOptions(
   firstLine: string,
   parseConfigPragmaFn: typeof ParseConfigPragma,
@@ -44,15 +48,11 @@ function makePluginOptions(
   ValueKindEnum: typeof ValueKind,
 ): [PluginOptions, Array<{filename: string | null; event: LoggerEvent}>] {
   let gating = null;
-  let enableEmitInstrumentForget = null;
-  let enableEmitFreeze = null;
-  let enableEmitHookGuards = null;
   let compilationMode: CompilationMode = 'all';
   let panicThreshold: PanicThresholdOptions = 'all_errors';
   let hookPattern: string | null = null;
   // TODO(@mofeiZ) rewrite snap fixtures to @validatePreserveExistingMemo:false
   let validatePreserveExistingMemoizationGuarantees = false;
-  let enableChangeDetectionForDebugging = null;
   let customMacros: null | Array<Macro> = null;
   let validateBlocklistedImports = null;
   let target = '19' as const;
@@ -78,31 +78,6 @@ function makePluginOptions(
       importSpecifierName: 'isForgetEnabled_Fixtures',
     };
   }
-  if (firstLine.includes('@instrumentForget')) {
-    enableEmitInstrumentForget = {
-      fn: {
-        source: 'react-compiler-runtime',
-        importSpecifierName: 'useRenderCounter',
-      },
-      gating: {
-        source: 'react-compiler-runtime',
-        importSpecifierName: 'shouldInstrument',
-      },
-      globalGating: '__DEV__',
-    };
-  }
-  if (firstLine.includes('@enableEmitFreeze')) {
-    enableEmitFreeze = {
-      source: 'react-compiler-runtime',
-      importSpecifierName: 'makeReadOnly',
-    };
-  }
-  if (firstLine.includes('@enableEmitHookGuards')) {
-    enableEmitHookGuards = {
-      source: 'react-compiler-runtime',
-      importSpecifierName: '$dispatcherGuard',
-    };
-  }
 
   const targetMatch = /@target="([^"]+)"/.exec(firstLine);
   if (targetMatch) {
@@ -132,16 +107,18 @@ function makePluginOptions(
     ignoreUseNoForget = true;
   }
 
+  /**
+   * Snap currently runs all fixtures without `validatePreserveExistingMemo` as
+   * most fixtures are interested in compilation output, not whether the
+   * compiler was able to preserve existing memo.
+   *
+   * TODO: flip the default. `useMemo` is rare in test fixtures -- fixtures that
+   * use useMemo should be explicit about whether this flag is enabled
+   */
   if (firstLine.includes('@validatePreserveExistingMemoizationGuarantees')) {
     validatePreserveExistingMemoizationGuarantees = true;
   }
 
-  if (firstLine.includes('@enableChangeDetectionForDebugging')) {
-    enableChangeDetectionForDebugging = {
-      source: 'react-compiler-runtime',
-      importSpecifierName: '$structuralCheck',
-    };
-  }
   const hookPatternMatch = /@hookPattern:"([^"]+)"/.exec(firstLine);
   if (
     hookPatternMatch &&
@@ -197,22 +174,6 @@ function makePluginOptions(
       .filter(s => s.length > 0);
   }
 
-  let lowerContextAccess = null;
-  if (firstLine.includes('@lowerContextAccess')) {
-    lowerContextAccess = {
-      source: 'react-compiler-runtime',
-      importSpecifierName: 'useContext_withSelector',
-    };
-  }
-
-  let inlineJsxTransform: EnvironmentConfig['inlineJsxTransform'] = null;
-  if (firstLine.includes('@enableInlineJsxTransform')) {
-    inlineJsxTransform = {
-      elementSymbol: 'react.transitional.element',
-      globalDevVar: 'DEV',
-    };
-  }
-
   let logs: Array<{filename: string | null; event: LoggerEvent}> = [];
   let logger: Logger | null = null;
   if (firstLine.includes('@logger')) {
@@ -232,17 +193,10 @@ function makePluginOptions(
         ValueKindEnum,
       }),
       customMacros,
-      enableEmitFreeze,
-      enableEmitInstrumentForget,
-      enableEmitHookGuards,
       assertValidMutableRanges: true,
-      enableSharedRuntime__testonly: true,
       hookPattern,
       validatePreserveExistingMemoizationGuarantees,
-      enableChangeDetectionForDebugging,
-      lowerContextAccess,
       validateBlocklistedImports,
-      inlineJsxTransform,
     },
     compilationMode,
     logger,