JSX Outlining

Adds a new pass to outline nested jsx to a separate function.

For now, we only outline nested jsx inside callbacks, but this can be
extended in the future (for ex, to outline nested jsx inside loops).

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

@@ -103,6 +103,7 @@ import {lowerContextAccess} from '../Optimization/LowerContextAccess';
 import {validateNoSetStateInPassiveEffects} from '../Validation/ValidateNoSetStateInPassiveEffects';
 import {validateNoJSXInTryStatement} from '../Validation/ValidateNoJSXInTryStatement';
 import {propagateScopeDependenciesHIR} from '../HIR/PropagateScopeDependenciesHIR';
+import {outlineJSX} from '../Optimization/OutlineJsx';
 
 export type CompilerPipelineValue =
  | {kind: 'ast'; name: string; value: CodegenFunction}
@@ -278,6 +279,10 @@ function* runWithEnvironment(
  value: hir,
  });
 
+ if (env.config.enableJsxOutlining) {
+ outlineJSX(hir);
+ }
+
  if (env.config.enableFunctionOutlining) {
  outlineFunctions(hir, fbtOperands);
  yield log({kind: 'hir', name: 'OutlineFunctions', value: hir});

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

@@ -199,7 +199,7 @@ function insertNewOutlinedFunctionNode(
  program: NodePath<t.Program>,
  originalFn: BabelFn,
  compiledFn: CodegenFunction,
-): NodePath<t.Function> {
+): BabelFn {
  switch (originalFn.type) {
  case 'FunctionDeclaration': {
  return originalFn.insertAfter(
@@ -491,18 +491,11 @@ export function compileProgram(
  fn.skip();
  ALREADY_COMPILED.add(fn.node);
  if (outlined.type !== null) {
- CompilerError.throwTodo({
- reason: `Implement support for outlining React functions (components/hooks)`,
- loc: outlined.fn.loc,
+ queue.push({
+ kind: 'outlined',
+ fn,
+ fnType: outlined.type,
  });
- /*
- * Above should be as simple as the following, but needs testing:
- * queue.push({
- * kind: "outlined",
- * fn,
- * fnType: outlined.type,
- * });
- */
  }
  }
  compiledFns.push({

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

@@ -354,6 +354,54 @@ const EnvironmentConfigSchema = z.object({
  */
  enableFunctionOutlining: z.boolean().default(true),
 
+ /**
+ * If enabled, this will outline nested JSX into a separate component.
+ *
+ * This will enable the compiler to memoize the separate component, giving us
+ * the same behavior as compiling _within_ the callback.
+ *
+ * ```
+ * function Component(countries, onDelete) {
+ * const name = useFoo();
+ * return countries.map(() => {
+ * return (
+ * <Foo>
+ * <Bar>{name}</Bar>
+ * <Button onclick={onDelete}>delete</Button>
+ * </Foo>
+ * );
+ * });
+ * }
+ * ```
+ *
+ * will be transpiled to:
+ *
+ * ```
+ * function Component(countries, onDelete) {
+ * const name = useFoo();
+ * return countries.map(() => {
+ * return (
+ * <Temp name={name} onDelete={onDelete} />
+ * );
+ * });
+ * }
+ *
+ * function Temp({name, onDelete}) {
+ * return (
+ * <Foo>
+ * <Bar>{name}</Bar>
+ * <Button onclick={onDelete}>delete</Button>
+ * </Foo>
+ * );
+ * }
+ *
+ * Both, `Component` and `Temp` will then be memoized by the compiler.
+ *
+ * With this change, when `countries` is updated by adding one single value,
+ * only the newly added value is re-rendered and not the entire list.
+ */
+ enableJsxOutlining: z.boolean().default(false),
+
  /*
  * Enables instrumentation codegen. This emits a dev-mode only call to an
  * instrumentation function, for components and hooks that Forget compiles.

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

@@ -920,15 +920,7 @@ export type InstructionValue =
  type: Type;
  loc: SourceLocation;
  }
- | {
- kind: 'JsxExpression';
- tag: Place | BuiltinTag;
- props: Array<JsxAttribute>;
- children: Array<Place> | null; // null === no children
- loc: SourceLocation;
- openingLoc: SourceLocation;
- closingLoc: SourceLocation;
- }
+ | JsxExpression
  | {
  kind: 'ObjectExpression';
  properties: Array<ObjectProperty | SpreadPattern>;
@@ -1074,6 +1066,16 @@ export type InstructionValue =
  loc: SourceLocation;
  };
 
+export type JsxExpression = {
+ kind: 'JsxExpression';
+ tag: Place | BuiltinTag;
+ props: Array<JsxAttribute>;
+ children: Array<Place> | null; // null === no children
+ loc: SourceLocation;
+ openingLoc: SourceLocation;
+ closingLoc: SourceLocation;
+};
+
 export type JsxAttribute =
  | {kind: 'JsxSpreadAttribute'; argument: Place}
  | {kind: 'JsxAttribute'; name: string; place: Place};