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).
facebook · Sep 16, 2024 · 5feb8d8 · 5feb8d8
1 parent 206df66
commit 5feb8d8
Show file tree

Hide file tree

Showing 8 changed files with 498 additions and 21 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
@@ -102,6 +102,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}
@@ -277,6 +278,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});

diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts b/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(
@@ -492,18 +492,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({

diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts
@@ -338,6 +338,11 @@ const EnvironmentConfigSchema = z.object({
    */
   enableFunctionOutlining: z.boolean().default(true),
 
+  /**
+   * TODO(gsn): Fill this out
+   */
+  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.

diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts
@@ -921,15 +921,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>;
@@ -1075,6 +1067,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};