[nocommit] Sketch of InlineJsxTransform

[ghstack-poisoned]

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts

@@ -106,6 +106,7 @@ import {propagatePhiTypes} from '../TypeInference/PropagatePhiTypes';
 import {lowerContextAccess} from '../Optimization/LowerContextAccess';
 import {validateNoSetStateInPassiveEffects} from '../Validation/ValidateNoSetStateInPassiveEffects';
 import {validateNoJSXInTryStatement} from '../Validation/ValidateNoJSXInTryStatement';
+import {inlineJsxTransform} from '../Optimization/InlineJSXTransform';
 
 export type CompilerPipelineValue =
   | {kind: 'ast'; name: string; value: CodegenFunction}
@@ -349,6 +350,21 @@ function* runWithEnvironment(
     assertTerminalPredsExist(hir);
   }
 
+  if (
+    env.config.enableInlineJsxTransform &&
+    // adding multiple instructions (after creating scopes)
+    // is only safe when enableReactiveScopesInHIR, where we
+    // can update the scope ranges after modifying instructions
+    env.config.enableReactiveScopesInHIR
+  ) {
+    inlineJsxTransform(hir);
+    yield log({
+      kind: 'hir',
+      name: 'InlineJsxTransform',
+      value: hir,
+    });
+  }
+
   const reactiveFunction = buildReactiveFunction(hir);
   yield log({
     kind: 'reactive',

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

@@ -232,6 +232,8 @@ const EnvironmentConfigSchema = z.object({
    */
   enableOptionalDependencies: z.boolean().default(true),
 
+  enableInlineJsxTransform: z.boolean().default(false),
+
   /*
    * Enable validation of hooks to partially check that the component honors the rules of hooks.
    * When disabled, the component is assumed to follow the rules (though the Babel plugin looks

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

@@ -0,0 +1,79 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import {
+  HIRFunction,
+  Instruction,
+  makeIdentifierId,
+  makeInstructionId,
+} from '../HIR';
+import {createTemporaryPlace, markInstructionIds} from '../HIR/HIRBuilder';
+
+export function inlineJsxTransform(fn: HIRFunction): void {
+  for (const [, block] of fn.body.blocks) {
+    let nextInstructions: Array<Instruction> | null = null;
+    for (let i = 0; i < block.instructions.length; i++) {
+      const instr = block.instructions[i]!;
+      switch (instr.value.kind) {
+        case 'JsxExpression': {
+          nextInstructions ??= block.instructions.slice(0, i);
+          const innerObject = createTemporaryPlace(fn.env, instr.value.loc);
+          nextInstructions.push({
+            id: makeInstructionId(0),
+            lvalue: {...innerObject},
+            value: {
+              kind: 'ObjectExpression',
+              properties: [],
+              loc: instr.value.loc,
+            },
+            loc: instr.loc,
+          });
+          const outerObject = createTemporaryPlace(fn.env, instr.value.loc);
+          nextInstructions.push({
+            id: makeInstructionId(0),
+            lvalue: {...outerObject},
+            value: {
+              kind: 'ObjectExpression',
+              properties: [
+                {
+                  kind: 'ObjectProperty',
+                  key: {kind: 'string', name: 'inner'},
+                  type: 'property',
+                  place: {...innerObject},
+                },
+              ],
+              loc: instr.value.loc,
+            },
+            loc: instr.loc,
+          });
+          // above creates
+          // {
+          //   inner: {}
+          // }
+          // don't emit the `instr`, create new instructions and push to nextInstructions
+          break;
+        }
+        case 'JsxFragment': {
+          nextInstructions ??= block.instructions.slice(0, i);
+          // similar to above
+          // don't emit the `instr`, create new instructions and push to nextInstructions
+          break;
+        }
+        default: {
+          if (nextInstructions !== null) {
+            nextInstructions.push(instr);
+          }
+        }
+      }
+    }
+    if (nextInstructions !== null) {
+      block.instructions = nextInstructions;
+    }
+  }
+  // TODO: fixup instruction ids, mutable ranges, etc
+  markInstructionIds(fn.body);
+}

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

@@ -0,0 +1 @@
+// @enableInlineJsxTransform