From 260f27301ab98eb2d8b57641b5b8127cb57f1052 Mon Sep 17 00:00:00 2001
From: Jack Pope <jackpope@meta.com>
Date: Mon, 16 Sep 2024 17:38:42 -0400
Subject: [PATCH] Add element symbol to config

---
 .../src/Entrypoint/Pipeline.ts                  |  4 ++--
 .../src/HIR/Environment.ts                      |  9 ++++++++-
 .../src/Optimization/InlineJsxTransform.ts      | 17 ++++++-----------
 compiler/packages/snap/src/compiler.ts          |  7 +++++++
 4 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts
index 606440b241f9b..900e4f4908494 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts
@@ -352,8 +352,8 @@ function* runWithEnvironment(
     });
   }
 
-  if (env.config.enableInlineJsxTransform) {
-    inlineJsxTransform(hir);
+  if (env.config.inlineJsxTransform) {
+    inlineJsxTransform(hir, env.config.inlineJsxTransform);
     yield log({
       kind: 'hir',
       name: 'inlineJsxTransform',
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 66270345fdf35..1af1df940b3fd 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts
@@ -50,6 +50,11 @@ import {
 import {Scope as BabelScope} from '@babel/traverse';
 import {TypeSchema} from './TypeSchema';
 
+export const ReactElementSymbolSchema = z.union([
+  z.literal('react.element'),
+  z.literal('react.transitional.element'),
+]);
+
 export const ExternalFunctionSchema = z.object({
   // Source for the imported module that exports the `importSpecifierName` functions
   source: z.string(),
@@ -237,8 +242,10 @@ const EnvironmentConfigSchema = z.object({
    * Enables inlining ReactElement object literals in place of JSX
    * An alternative to the standard JSX transform which replaces JSX with React's jsxProd() runtime
    * Currently a prod-only optimization, requiring Fast JSX dependencies
+   *
+   * The symbol configuration is set for backwards compatability with pre-React 19 transforms
    */
-  enableInlineJsxTransform: z.boolean().default(false),
+  inlineJsxTransform: ReactElementSymbolSchema.nullish(),
 
   /*
    * Enable validation of hooks to partially check that the component honors the rules of hooks.
diff --git a/compiler/packages/babel-plugin-react-compiler/src/Optimization/InlineJsxTransform.ts b/compiler/packages/babel-plugin-react-compiler/src/Optimization/InlineJsxTransform.ts
index 0fe516da977e3..580ef5e7bc0a4 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Optimization/InlineJsxTransform.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Optimization/InlineJsxTransform.ts
@@ -316,7 +316,10 @@ function createPropsProperties(
 }
 
 // TODO: Make PROD only with conditional statements
-export function inlineJsxTransform(fn: HIRFunction): void {
+export function inlineJsxTransform(
+  fn: HIRFunction,
+  elementSymbol: 'react.element' | 'react.transitional.element',
+): void {
   for (const [, block] of fn.body.blocks) {
     let nextInstructions: Array<Instruction> | null = null;
     for (let i = 0; i < block.instructions.length; i++) {
@@ -344,11 +347,7 @@ export function inlineJsxTransform(fn: HIRFunction): void {
                   instr,
                   nextInstructions,
                   '$$typeof',
-                  /**
-                   * TODO: Add this to config so we can switch between
-                   * react.element / react.transitional.element
-                   */
-                  'react.transitional.element',
+                  elementSymbol,
                 ),
                 createTagProperty(fn, instr, nextInstructions, instr.value.tag),
                 refProperty,
@@ -384,11 +383,7 @@ export function inlineJsxTransform(fn: HIRFunction): void {
                   instr,
                   nextInstructions,
                   '$$typeof',
-                  /**
-                   * TODO: Add this to config so we can switch between
-                   * react.element / react.transitional.element
-                   */
-                  'react.transitional.element',
+                  elementSymbol,
                 ),
                 createSymbolProperty(
                   fn,
diff --git a/compiler/packages/snap/src/compiler.ts b/compiler/packages/snap/src/compiler.ts
index 417a657d28012..3587d1411eb77 100644
--- a/compiler/packages/snap/src/compiler.ts
+++ b/compiler/packages/snap/src/compiler.ts
@@ -21,6 +21,7 @@ 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,
@@ -201,6 +202,11 @@ function makePluginOptions(
     };
   }
 
+  let inlineJsxTransform: EnvironmentConfig['inlineJsxTransform'] = null;
+  if (firstLine.includes('@enableInlineJsxTransform')) {
+    inlineJsxTransform = 'react.transitional.element';
+  }
+
   let logs: Array<{filename: string | null; event: LoggerEvent}> = [];
   let logger: Logger | null = null;
   if (firstLine.includes('@logger')) {
@@ -230,6 +236,7 @@ function makePluginOptions(
       enableChangeDetectionForDebugging,
       lowerContextAccess,
       validateBlocklistedImports,
+      inlineJsxTransform,
     },
     compilationMode,
     logger,