[compiler] Wrap inline jsx transform codegen in conditional (#31267)

JSX inlining is a prod-only optimization. We want to enforce this while
maintaining the same compiler output in DEV and PROD.

Here we add a conditional to the transform that only replaces JSX with
object literals outside of DEV. Then a later build step can handle DCE
based on the value of `__DEV__`

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

@@ -55,6 +55,7 @@ export const ReactElementSymbolSchema = z.object({
     z.literal('react.element'),
     z.literal('react.transitional.element'),
   ]),
+  globalDevVar: z.string(),
 });
 
 export const ExternalFunctionSchema = z.object({

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

@@ -1243,6 +1243,17 @@ export function makeTemporaryIdentifier(
   };
 }
 
+export function forkTemporaryIdentifier(
+  id: IdentifierId,
+  source: Identifier,
+): Identifier {
+  return {
+    ...source,
+    mutableRange: {start: makeInstructionId(0), end: makeInstructionId(0)},
+    id,
+  };
+}
+
 /**
  * Creates a valid identifier name. This should *not* be used for synthesizing
  * identifier names: only call this method for identifier names that appear in the

compiler/packages/babel-plugin-react-compiler/src/Optimization/InlineJsxTransform.ts

@@ -6,14 +6,25 @@
  */
 
 import {
+  BasicBlock,
+  BlockId,
   BuiltinTag,
+  DeclarationId,
   Effect,
+  forkTemporaryIdentifier,
+  GotoTerminal,
+  GotoVariant,
   HIRFunction,
+  Identifier,
+  IfTerminal,
   Instruction,
+  InstructionKind,
   JsxAttribute,
   makeInstructionId,
   ObjectProperty,
+  Phi,
   Place,
+  promoteTemporary,
   SpreadPattern,
 } from '../HIR';
 import {
@@ -24,6 +35,365 @@ import {
   reversePostorderBlocks,
 } from '../HIR/HIRBuilder';
 import {CompilerError, EnvironmentConfig} from '..';
+import {
+  mapInstructionLValues,
+  mapInstructionOperands,
+  mapInstructionValueOperands,
+  mapTerminalOperands,
+} from '../HIR/visitors';
+
+type InlinedJsxDeclarationMap = Map<
+  DeclarationId,
+  {identifier: Identifier; blockIdsToIgnore: Set<BlockId>}
+>;
+
+/**
+ * A prod-only, RN optimization to replace JSX with inlined ReactElement object literals
+ *
+ * Example:
+ * <>foo</>
+ * _______________
+ * let t1;
+ * if (__DEV__) {
+ *   t1 = <>foo</>
+ * } else {
+ *   t1 = {...}
+ * }
+ *
+ */
+export function inlineJsxTransform(
+  fn: HIRFunction,
+  inlineJsxTransformConfig: NonNullable<
+    EnvironmentConfig['inlineJsxTransform']
+  >,
+): void {
+  const inlinedJsxDeclarations: InlinedJsxDeclarationMap = new Map();
+  /**
+   * Step 1: Codegen the conditional and ReactElement object literal
+   */
+  for (const [_, currentBlock] of [...fn.body.blocks]) {
+    let fallthroughBlockInstructions: Array<Instruction> | null = null;
+    const instructionCount = currentBlock.instructions.length;
+    for (let i = 0; i < instructionCount; i++) {
+      const instr = currentBlock.instructions[i]!;
+      // TODO: Support value blocks
+      if (currentBlock.kind === 'value') {
+        fn.env.logger?.logEvent(fn.env.filename, {
+          kind: 'CompileDiagnostic',
+          fnLoc: null,
+          detail: {
+            reason: 'JSX Inlining is not supported on value blocks',
+            loc: instr.loc,
+          },
+        });
+        continue;
+      }
+      switch (instr.value.kind) {
+        case 'JsxExpression':
+        case 'JsxFragment': {
+          /**
+           * Split into blocks for new IfTerminal:
+           *   current, then, else, fallthrough
+           */
+          const currentBlockInstructions = currentBlock.instructions.slice(
+            0,
+            i,
+          );
+          const thenBlockInstructions = currentBlock.instructions.slice(
+            i,
+            i + 1,
+          );
+          const elseBlockInstructions: Array<Instruction> = [];
+          fallthroughBlockInstructions ??= currentBlock.instructions.slice(
+            i + 1,
+          );
+
+          const fallthroughBlockId = fn.env.nextBlockId;
+          const fallthroughBlock: BasicBlock = {
+            kind: currentBlock.kind,
+            id: fallthroughBlockId,
+            instructions: fallthroughBlockInstructions,
+            terminal: currentBlock.terminal,
+            preds: new Set(),
+            phis: new Set(),
+          };
+
+          /**
+           * Complete current block
+           * - Add instruction for variable declaration
+           * - Add instruction for LoadGlobal used by conditional
+           * - End block with a new IfTerminal
+           */
+          const varPlace = createTemporaryPlace(fn.env, instr.value.loc);
+          promoteTemporary(varPlace.identifier);
+          const varLValuePlace = createTemporaryPlace(fn.env, instr.value.loc);
+          const thenVarPlace = {
+            ...varPlace,
+            identifier: forkTemporaryIdentifier(
+              fn.env.nextIdentifierId,
+              varPlace.identifier,
+            ),
+          };
+          const elseVarPlace = {
+            ...varPlace,
+            identifier: forkTemporaryIdentifier(
+              fn.env.nextIdentifierId,
+              varPlace.identifier,
+            ),
+          };
+          const varInstruction: Instruction = {
+            id: makeInstructionId(0),
+            lvalue: {...varLValuePlace},
+            value: {
+              kind: 'DeclareLocal',
+              lvalue: {place: {...varPlace}, kind: InstructionKind.Let},
+              type: null,
+              loc: instr.value.loc,
+            },
+            loc: instr.loc,
+          };
+          currentBlockInstructions.push(varInstruction);
+
+          const devGlobalPlace = createTemporaryPlace(fn.env, instr.value.loc);
+          const devGlobalInstruction: Instruction = {
+            id: makeInstructionId(0),
+            lvalue: {...devGlobalPlace, effect: Effect.Mutate},
+            value: {
+              kind: 'LoadGlobal',
+              binding: {
+                kind: 'Global',
+                name: inlineJsxTransformConfig.globalDevVar,
+              },
+              loc: instr.value.loc,
+            },
+            loc: instr.loc,
+          };
+          currentBlockInstructions.push(devGlobalInstruction);
+          const thenBlockId = fn.env.nextBlockId;
+          const elseBlockId = fn.env.nextBlockId;
+          const ifTerminal: IfTerminal = {
+            kind: 'if',
+            test: {...devGlobalPlace, effect: Effect.Read},
+            consequent: thenBlockId,
+            alternate: elseBlockId,
+            fallthrough: fallthroughBlockId,
+            loc: instr.loc,
+            id: makeInstructionId(0),
+          };
+          currentBlock.instructions = currentBlockInstructions;
+          currentBlock.terminal = ifTerminal;
+
+          /**
+           * Set up then block where we put the original JSX return
+           */
+          const thenBlock: BasicBlock = {
+            id: thenBlockId,
+            instructions: thenBlockInstructions,
+            kind: 'block',
+            phis: new Set(),
+            preds: new Set(),
+            terminal: {
+              kind: 'goto',
+              block: fallthroughBlockId,
+              variant: GotoVariant.Break,
+              id: makeInstructionId(0),
+              loc: instr.loc,
+            },
+          };
+          fn.body.blocks.set(thenBlockId, thenBlock);
+
+          const resassignElsePlace = createTemporaryPlace(
+            fn.env,
+            instr.value.loc,
+          );
+          const reassignElseInstruction: Instruction = {
+            id: makeInstructionId(0),
+            lvalue: {...resassignElsePlace},
+            value: {
+              kind: 'StoreLocal',
+              lvalue: {
+                place: elseVarPlace,
+                kind: InstructionKind.Reassign,
+              },
+              value: {...instr.lvalue},
+              type: null,
+              loc: instr.value.loc,
+            },
+            loc: instr.loc,
+          };
+          thenBlockInstructions.push(reassignElseInstruction);
+
+          /**
+           * Set up else block where we add new codegen
+           */
+          const elseBlockTerminal: GotoTerminal = {
+            kind: 'goto',
+            block: fallthroughBlockId,
+            variant: GotoVariant.Break,
+            id: makeInstructionId(0),
+            loc: instr.loc,
+          };
+          const elseBlock: BasicBlock = {
+            id: elseBlockId,
+            instructions: elseBlockInstructions,
+            kind: 'block',
+            phis: new Set(),
+            preds: new Set(),
+            terminal: elseBlockTerminal,
+          };
+          fn.body.blocks.set(elseBlockId, elseBlock);
+
+          /**
+           * ReactElement object literal codegen
+           */
+          const {refProperty, keyProperty, propsProperty} =
+            createPropsProperties(
+              fn,
+              instr,
+              elseBlockInstructions,
+              instr.value.kind === 'JsxExpression' ? instr.value.props : [],
+              instr.value.children,
+            );
+          const reactElementInstructionPlace = createTemporaryPlace(
+            fn.env,
+            instr.value.loc,
+          );
+          const reactElementInstruction: Instruction = {
+            id: makeInstructionId(0),
+            lvalue: {...reactElementInstructionPlace, effect: Effect.Store},
+            value: {
+              kind: 'ObjectExpression',
+              properties: [
+                createSymbolProperty(
+                  fn,
+                  instr,
+                  elseBlockInstructions,
+                  '$$typeof',
+                  inlineJsxTransformConfig.elementSymbol,
+                ),
+                instr.value.kind === 'JsxExpression'
+                  ? createTagProperty(
+                      fn,
+                      instr,
+                      elseBlockInstructions,
+                      instr.value.tag,
+                    )
+                  : createSymbolProperty(
+                      fn,
+                      instr,
+                      elseBlockInstructions,
+                      'type',
+                      'react.fragment',
+                    ),
+                refProperty,
+                keyProperty,
+                propsProperty,
+              ],
+              loc: instr.value.loc,
+            },
+            loc: instr.loc,
+          };
+          elseBlockInstructions.push(reactElementInstruction);
+
+          const reassignConditionalInstruction: Instruction = {
+            id: makeInstructionId(0),
+            lvalue: {...createTemporaryPlace(fn.env, instr.value.loc)},
+            value: {
+              kind: 'StoreLocal',
+              lvalue: {
+                place: {...elseVarPlace},
+                kind: InstructionKind.Reassign,
+              },
+              value: {...reactElementInstruction.lvalue},
+              type: null,
+              loc: instr.value.loc,
+            },
+            loc: instr.loc,
+          };
+          elseBlockInstructions.push(reassignConditionalInstruction);
+
+          /**
+           * Create phis to reassign the var
+           */
+          const operands: Map<BlockId, Place> = new Map();
+          operands.set(thenBlockId, {
+            ...elseVarPlace,
+          });
+          operands.set(elseBlockId, {
+            ...thenVarPlace,
+          });
+
+          const phiIdentifier = forkTemporaryIdentifier(
+            fn.env.nextIdentifierId,
+            varPlace.identifier,
+          );
+          const phiPlace = {
+            ...createTemporaryPlace(fn.env, instr.value.loc),
+            identifier: phiIdentifier,
+          };
+          const phis: Set<Phi> = new Set([
+            {
+              kind: 'Phi',
+              operands,
+              place: phiPlace,
+            },
+          ]);
+          fallthroughBlock.phis = phis;
+          fn.body.blocks.set(fallthroughBlockId, fallthroughBlock);
+
+          /**
+           * Track this JSX instruction so we can replace references in step 2
+           */
+          inlinedJsxDeclarations.set(instr.lvalue.identifier.declarationId, {
+            identifier: phiIdentifier,
+            blockIdsToIgnore: new Set([thenBlockId, elseBlockId]),
+          });
+          break;
+        }
+        case 'FunctionExpression':
+        case 'ObjectMethod': {
+          inlineJsxTransform(
+            instr.value.loweredFunc.func,
+            inlineJsxTransformConfig,
+          );
+          break;
+        }
+      }
+    }
+  }
+
+  /**
+   * Step 2: Replace declarations with new phi values
+   */
+  for (const [blockId, block] of fn.body.blocks) {
+    for (const instr of block.instructions) {
+      mapInstructionOperands(instr, place =>
+        handlePlace(place, blockId, inlinedJsxDeclarations),
+      );
+
+      mapInstructionLValues(instr, lvalue =>
+        handlelValue(lvalue, blockId, inlinedJsxDeclarations),
+      );
+
+      mapInstructionValueOperands(instr.value, place =>
+        handlePlace(place, blockId, inlinedJsxDeclarations),
+      );
+    }
+
+    mapTerminalOperands(block.terminal, place =>
+      handlePlace(place, blockId, inlinedJsxDeclarations),
+    );
+  }
+
+  /**
+   * Step 3: Fixup the HIR
+   * Restore RPO, ensure correct predecessors, renumber instructions, fix scope and ranges.
+   */
+  reversePostorderBlocks(fn.body);
+  markPredecessors(fn.body);
+  markInstructionIds(fn.body);
+  fixScopeAndIdentifierRanges(fn.body);
+}
 
 function createSymbolProperty(
   fn: HIRFunction,
@@ -315,123 +685,38 @@ function createPropsProperties(
   return {refProperty, keyProperty, propsProperty};
 }
 
-// TODO: Make PROD only with conditional statements
+function handlePlace(
-export function inlineJsxTransform(
+  place: Place,
-  fn: HIRFunction,
+  blockId: BlockId,
-  inlineJsxTransformConfig: NonNullable<
+  inlinedJsxDeclarations: InlinedJsxDeclarationMap,
-    EnvironmentConfig['inlineJsxTransform']
+): Place {
-  >,
+  const inlinedJsxDeclaration = inlinedJsxDeclarations.get(
-): void {
+    place.identifier.declarationId,
-  for (const [, block] of fn.body.blocks) {
+  );
-    let nextInstructions: Array<Instruction> | null = null;
+  if (
-    for (let i = 0; i < block.instructions.length; i++) {
+    inlinedJsxDeclaration == null ||
-      const instr = block.instructions[i]!;
+    inlinedJsxDeclaration.blockIdsToIgnore.has(blockId)
-      switch (instr.value.kind) {
+  ) {
-        case 'JsxExpression': {
+    return {...place};
-          nextInstructions ??= block.instructions.slice(0, i);
-
-          const {refProperty, keyProperty, propsProperty} =
-            createPropsProperties(
-              fn,
-              instr,
-              nextInstructions,
-              instr.value.props,
-              instr.value.children,
-            );
-          const reactElementInstruction: Instruction = {
-            id: makeInstructionId(0),
-            lvalue: {...instr.lvalue, effect: Effect.Store},
-            value: {
-              kind: 'ObjectExpression',
-              properties: [
-                createSymbolProperty(
-                  fn,
-                  instr,
-                  nextInstructions,
-                  '$$typeof',
-                  inlineJsxTransformConfig.elementSymbol,
-                ),
-                createTagProperty(fn, instr, nextInstructions, instr.value.tag),
-                refProperty,
-                keyProperty,
-                propsProperty,
-              ],
-              loc: instr.value.loc,
-            },
-            loc: instr.loc,
-          };
-          nextInstructions.push(reactElementInstruction);
-
-          break;
-        }
-        case 'JsxFragment': {
-          nextInstructions ??= block.instructions.slice(0, i);
-          const {refProperty, keyProperty, propsProperty} =
-            createPropsProperties(
-              fn,
-              instr,
-              nextInstructions,
-              [],
-              instr.value.children,
-            );
-          const reactElementInstruction: Instruction = {
-            id: makeInstructionId(0),
-            lvalue: {...instr.lvalue, effect: Effect.Store},
-            value: {
-              kind: 'ObjectExpression',
-              properties: [
-                createSymbolProperty(
-                  fn,
-                  instr,
-                  nextInstructions,
-                  '$$typeof',
-                  inlineJsxTransformConfig.elementSymbol,
-                ),
-                createSymbolProperty(
-                  fn,
-                  instr,
-                  nextInstructions,
-                  'type',
-                  'react.fragment',
-                ),
-                refProperty,
-                keyProperty,
-                propsProperty,
-              ],
-              loc: instr.value.loc,
-            },
-            loc: instr.loc,
-          };
-          nextInstructions.push(reactElementInstruction);
-          break;
-        }
-        case 'FunctionExpression':
-        case 'ObjectMethod': {
-          inlineJsxTransform(
-            instr.value.loweredFunc.func,
-            inlineJsxTransformConfig,
-          );
-          if (nextInstructions !== null) {
-            nextInstructions.push(instr);
-          }
-          break;
-        }
-        default: {
-          if (nextInstructions !== null) {
-            nextInstructions.push(instr);
-          }
-        }
-      }
-    }
-    if (nextInstructions !== null) {
-      block.instructions = nextInstructions;
-    }
   }
 
-  // Fixup the HIR to restore RPO, ensure correct predecessors, and renumber instructions.
+  return {...place, identifier: {...inlinedJsxDeclaration.identifier}};
-  reversePostorderBlocks(fn.body);
+}
-  markPredecessors(fn.body);
+
-  markInstructionIds(fn.body);
+function handlelValue(
-  // The renumbering instructions invalidates scope and identifier ranges
+  lvalue: Place,
-  fixScopeAndIdentifierRanges(fn.body);
+  blockId: BlockId,
+  inlinedJsxDeclarations: InlinedJsxDeclarationMap,
+): Place {
+  const inlinedJsxDeclaration = inlinedJsxDeclarations.get(
+    lvalue.identifier.declarationId,
+  );
+  if (
+    inlinedJsxDeclaration == null ||
+    inlinedJsxDeclaration.blockIdsToIgnore.has(blockId)
+  ) {
+    return {...lvalue};
+  }
+
+  return {...lvalue, identifier: {...inlinedJsxDeclaration.identifier}};
 }

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

@@ -50,6 +50,22 @@ function PropsSpread() {
   );
 }
 
+function ConditionalJsx({shouldWrap}) {
+  let content = <div>Hello</div>;
+
+  if (shouldWrap) {
+    content = <Parent>{content}</Parent>;
+  }
+
+  return content;
+}
+
+// TODO: Support value blocks
+function TernaryJsx({cond}) {
+  return cond ? <div /> : null;
+}
+
+global.DEV = true;
 export const FIXTURE_ENTRYPOINT = {
   fn: ParentAndChildren,
   params: [{foo: 'abc'}],
@@ -67,13 +83,17 @@ function Parent(t0) {
   const { children, ref } = t0;
   let t1;
   if ($[0] !== children) {
-    t1 = {
+    if (DEV) {
-      $$typeof: Symbol.for("react.transitional.element"),
+      t1 = <div ref={ref}>{children}</div>;
-      type: "div",
+    } else {
-      ref: ref,
+      t1 = {
-      key: null,
+        $$typeof: Symbol.for("react.transitional.element"),
-      props: { children: children },
+        type: "div",
-    };
+        ref: ref,
+        key: null,
+        props: { children: children },
+      };
+    }
     $[0] = children;
     $[1] = t1;
   } else {
@@ -87,13 +107,17 @@ function Child(t0) {
   const { children } = t0;
   let t1;
   if ($[0] !== children) {
-    t1 = {
+    if (DEV) {
-      $$typeof: Symbol.for("react.transitional.element"),
+      t1 = <>{children}</>;
-      type: Symbol.for("react.fragment"),
+    } else {
-      ref: null,
+      t1 = {
-      key: null,
+        $$typeof: Symbol.for("react.transitional.element"),
-      props: { children: children },
+        type: Symbol.for("react.fragment"),
-    };
+        ref: null,
+        key: null,
+        props: { children: children },
+      };
+    }
     $[0] = children;
     $[1] = t1;
   } else {
@@ -107,26 +131,34 @@ function GrandChild(t0) {
   const { className } = t0;
   let t1;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t1 = {
+    if (DEV) {
-      $$typeof: Symbol.for("react.transitional.element"),
+      t1 = <React.Fragment key="fragmentKey">Hello world</React.Fragment>;
-      type: React.Fragment,
+    } else {
-      ref: null,
+      t1 = {
-      key: "fragmentKey",
+        $$typeof: Symbol.for("react.transitional.element"),
-      props: { children: "Hello world" },
+        type: React.Fragment,
-    };
+        ref: null,
+        key: "fragmentKey",
+        props: { children: "Hello world" },
+      };
+    }
     $[0] = t1;
   } else {
     t1 = $[0];
   }
   let t2;
   if ($[1] !== className) {
-    t2 = {
+    if (DEV) {
-      $$typeof: Symbol.for("react.transitional.element"),
+      t2 = <span className={className}>{t1}</span>;
-      type: "span",
+    } else {
-      ref: null,
+      t2 = {
-      key: null,
+        $$typeof: Symbol.for("react.transitional.element"),
-      props: { className: className, children: t1 },
+        type: "span",
-    };
+        ref: null,
+        key: null,
+        props: { className: className, children: t1 },
+      };
+    }
     $[1] = className;
     $[2] = t2;
   } else {
@@ -140,13 +172,17 @@ function ParentAndRefAndKey(props) {
   const testRef = useRef();
   let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t0 = {
+    if (DEV) {
-      $$typeof: Symbol.for("react.transitional.element"),
+      t0 = <Parent a="a" b={{ b: "b" }} c={C} key="testKey" ref={testRef} />;
-      type: Parent,
+    } else {
-      ref: testRef,
+      t0 = {
-      key: "testKey",
+        $$typeof: Symbol.for("react.transitional.element"),
-      props: { a: "a", b: { b: "b" }, c: C },
+        type: Parent,
-    };
+        ref: testRef,
+        key: "testKey",
+        props: { a: "a", b: { b: "b" }, c: C },
+      };
+    }
     $[0] = t0;
   } else {
     t0 = $[0];
@@ -158,13 +194,21 @@ function ParentAndChildren(props) {
   const $ = _c2(14);
   let t0;
   if ($[0] !== props.foo) {
-    t0 = () => ({
+    t0 = () => {
-      $$typeof: Symbol.for("react.transitional.element"),
+      let t1;
-      type: "div",
+      if (DEV) {
-      ref: null,
+        t1 = <div key="d">{props.foo}</div>;
-      key: "d",
+      } else {
-      props: { children: props.foo },
+        t1 = {
-    });
+          $$typeof: Symbol.for("react.transitional.element"),
+          type: "div",
+          ref: null,
+          key: "d",
+          props: { children: props.foo },
+        };
+      }
+      return t1;
+    };
     $[0] = props.foo;
     $[1] = t0;
   } else {
@@ -173,71 +217,99 @@ function ParentAndChildren(props) {
   const render = t0;
   let t1;
   if ($[2] !== props) {
-    t1 = {
+    if (DEV) {
-      $$typeof: Symbol.for("react.transitional.element"),
+      t1 = <Child key="a" {...props} />;
-      type: Child,
+    } else {
-      ref: null,
+      t1 = {
-      key: "a",
+        $$typeof: Symbol.for("react.transitional.element"),
-      props: props,
+        type: Child,
-    };
+        ref: null,
+        key: "a",
+        props: props,
+      };
+    }
     $[2] = props;
     $[3] = t1;
   } else {
     t1 = $[3];
   }
-  let t2;
+
-  if ($[4] !== props) {
+  const t2 = props.foo;
-    t2 = {
-      $$typeof: Symbol.for("react.transitional.element"),
-      type: GrandChild,
-      ref: null,
-      key: "c",
-      props: { className: props.foo, ...props },
-    };
-    $[4] = props;
-    $[5] = t2;
-  } else {
-    t2 = $[5];
-  }
   let t3;
-  if ($[6] !== render) {
+  if ($[4] !== props) {
-    t3 = render();
+    if (DEV) {
-    $[6] = render;
+      t3 = <GrandChild key="c" className={t2} {...props} />;
-    $[7] = t3;
+    } else {
+      t3 = {
+        $$typeof: Symbol.for("react.transitional.element"),
+        type: GrandChild,
+        ref: null,
+        key: "c",
+        props: { className: t2, ...props },
+      };
+    }
+    $[4] = props;
+    $[5] = t3;
   } else {
-    t3 = $[7];
+    t3 = $[5];
   }
   let t4;
-  if ($[8] !== t2 || $[9] !== t3) {
+  if ($[6] !== render) {
-    t4 = {
+    t4 = render();
-      $$typeof: Symbol.for("react.transitional.element"),
+    $[6] = render;
-      type: Child,
+    $[7] = t4;
-      ref: null,
-      key: "b",
-      props: { children: [t2, t3] },
-    };
-    $[8] = t2;
-    $[9] = t3;
-    $[10] = t4;
   } else {
-    t4 = $[10];
+    t4 = $[7];
   }
   let t5;
-  if ($[11] !== t1 || $[12] !== t4) {
+  if ($[8] !== t3 || $[9] !== t4) {
-    t5 = {
+    if (DEV) {
-      $$typeof: Symbol.for("react.transitional.element"),
+      t5 = (
-      type: Parent,
+        <Child key="b">
-      ref: null,
+          {t3}
-      key: null,
+          {t4}
-      props: { children: [t1, t4] },
+        </Child>
-    };
+      );
-    $[11] = t1;
+    } else {
-    $[12] = t4;
+      t5 = {
-    $[13] = t5;
+        $$typeof: Symbol.for("react.transitional.element"),
+        type: Child,
+        ref: null,
+        key: "b",
+        props: { children: [t3, t4] },
+      };
+    }
+    $[8] = t3;
+    $[9] = t4;
+    $[10] = t5;
   } else {
-    t5 = $[13];
+    t5 = $[10];
   }
-  return t5;
+  let t6;
+  if ($[11] !== t1 || $[12] !== t5) {
+    if (DEV) {
+      t6 = (
+        <Parent>
+          {t1}
+          {t5}
+        </Parent>
+      );
+    } else {
+      t6 = {
+        $$typeof: Symbol.for("react.transitional.element"),
+        type: Parent,
+        ref: null,
+        key: null,
+        props: { children: [t1, t5] },
+      };
+    }
+    $[11] = t1;
+    $[12] = t5;
+    $[13] = t6;
+  } else {
+    t6 = $[13];
+  }
+  return t6;
 }
 
 const propsToSpread = { a: "a", b: "b", c: "c" };
@@ -245,30 +317,46 @@ function PropsSpread() {
   const $ = _c2(1);
   let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t0 = {
+    let t1;
-      $$typeof: Symbol.for("react.transitional.element"),
+    if (DEV) {
-      type: Symbol.for("react.fragment"),
+      t1 = <Test key="a" {...propsToSpread} />;
-      ref: null,
+    } else {
-      key: null,
+      t1 = {
-      props: {
+        $$typeof: Symbol.for("react.transitional.element"),
-        children: [
+        type: Test,
-          {
+        ref: null,
-            $$typeof: Symbol.for("react.transitional.element"),
+        key: "a",
-            type: Test,
+        props: propsToSpread,
-            ref: null,
+      };
-            key: "a",
+    }
-            props: propsToSpread,
+    let t2;
-          },
+    if (DEV) {
-          {
+      t2 = <Test key="b" {...propsToSpread} a="z" />;
-            $$typeof: Symbol.for("react.transitional.element"),
+    } else {
-            type: Test,
+      t2 = {
-            ref: null,
+        $$typeof: Symbol.for("react.transitional.element"),
-            key: "b",
+        type: Test,
-            props: { ...propsToSpread, a: "z" },
+        ref: null,
-          },
+        key: "b",
-        ],
+        props: { ...propsToSpread, a: "z" },
-      },
+      };
-    };
+    }
+    if (DEV) {
+      t0 = (
+        <>
+          {t1}
+          {t2}
+        </>
+      );
+    } else {
+      t0 = {
+        $$typeof: Symbol.for("react.transitional.element"),
+        type: Symbol.for("react.fragment"),
+        ref: null,
+        key: null,
+        props: { children: [t1, t2] },
+      };
+    }
     $[0] = t0;
   } else {
     t0 = $[0];
@@ -276,6 +364,67 @@ function PropsSpread() {
   return t0;
 }
 
+function ConditionalJsx(t0) {
+  const $ = _c2(2);
+  const { shouldWrap } = t0;
+  let t1;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    if (DEV) {
+      t1 = <div>Hello</div>;
+    } else {
+      t1 = {
+        $$typeof: Symbol.for("react.transitional.element"),
+        type: "div",
+        ref: null,
+        key: null,
+        props: { children: "Hello" },
+      };
+    }
+    $[0] = t1;
+  } else {
+    t1 = $[0];
+  }
+  let content = t1;
+  if (shouldWrap) {
+    const t2 = content;
+    let t3;
+    if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
+      if (DEV) {
+        t3 = <Parent>{t2}</Parent>;
+      } else {
+        t3 = {
+          $$typeof: Symbol.for("react.transitional.element"),
+          type: Parent,
+          ref: null,
+          key: null,
+          props: { children: t2 },
+        };
+      }
+      $[1] = t3;
+    } else {
+      t3 = $[1];
+    }
+    content = t3;
+  }
+  return content;
+}
+
+// TODO: Support value blocks
+function TernaryJsx(t0) {
+  const $ = _c2(2);
+  const { cond } = t0;
+  let t1;
+  if ($[0] !== cond) {
+    t1 = cond ? <div /> : null;
+    $[0] = cond;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  return t1;
+}
+
+global.DEV = true;
 export const FIXTURE_ENTRYPOINT = {
   fn: ParentAndChildren,
   params: [{ foo: "abc" }],

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

@@ -46,6 +46,22 @@ function PropsSpread() {
   );
 }
 
+function ConditionalJsx({shouldWrap}) {
+  let content = <div>Hello</div>;
+
+  if (shouldWrap) {
+    content = <Parent>{content}</Parent>;
+  }
+
+  return content;
+}
+
+// TODO: Support value blocks
+function TernaryJsx({cond}) {
+  return cond ? <div /> : null;
+}
+
+global.DEV = true;
 export const FIXTURE_ENTRYPOINT = {
   fn: ParentAndChildren,
   params: [{foo: 'abc'}],

compiler/packages/snap/src/compiler.ts

@@ -207,7 +207,10 @@ function makePluginOptions(
 
   let inlineJsxTransform: EnvironmentConfig['inlineJsxTransform'] = null;
   if (firstLine.includes('@enableInlineJsxTransform')) {
-    inlineJsxTransform = {elementSymbol: 'react.transitional.element'};
+    inlineJsxTransform = {
+      elementSymbol: 'react.transitional.element',
+      globalDevVar: 'DEV',
+    };
   }
 
   let logs: Array<{filename: string | null; event: LoggerEvent}> = [];