[compiler] Outline JSX with non-jsx children (#31442)

Previously, we bailed out on outlining jsx that had children that were
not part of the outlined jsx.

Now, we add support for children by treating as attributes.

Author: gsathya

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

@@ -219,10 +219,20 @@ type OutlinedJsxAttribute = {
 function collectProps(
   instructions: Array<JsxInstruction>,
 ): Array<OutlinedJsxAttribute> | null {
+  let id = 1;
+
+  function generateName(oldName: string): string {
+    let newName = oldName;
+    while (seen.has(newName)) {
+      newName = `${oldName}${id++}`;
+    }
+    seen.add(newName);
+    return newName;
+  }
+
   const attributes: Array<OutlinedJsxAttribute> = [];
   const jsxIds = new Set(instructions.map(i => i.lvalue.identifier.id));
   const seen: Set<string> = new Set();
-  let id = 1;
 
   for (const instr of instructions) {
     const {value} = instr;
@@ -233,25 +243,29 @@ function collectProps(
       }
 
       if (at.kind === 'JsxAttribute') {
-        let newName = at.name;
-        while (seen.has(newName)) {
-          newName = `${at.name}${id++}`;
-        }
+        const newName = generateName(at.name);
         attributes.push({
           originalName: at.name,
           newName,
           place: at.place,
         });
-        seen.add(newName);
       }
     }
 
-    // TODO(gsn): Add support for children that are not jsx expressions
-    if (
-      value.children &&
-      value.children.some(child => !jsxIds.has(child.identifier.id))
-    ) {
-      return null;
+    if (value.children) {
+      for (const child of value.children) {
+        if (jsxIds.has(child.identifier.id)) {
+          continue;
+        }
+
+        promoteTemporary(child.identifier);
+        const newName = generateName('t');
+        attributes.push({
+          originalName: child.identifier.name!.value,
+          newName: newName,
+          place: child,
+        });
+      }
     }
   }
   return attributes;
@@ -387,6 +401,7 @@ function emitUpdatedJsx(
   oldToNewProps: Map<IdentifierId, OutlinedJsxAttribute>,
 ): Array<JsxInstruction> {
   const newInstrs: Array<JsxInstruction> = [];
+  const jsxIds = new Set(jsx.map(i => i.lvalue.identifier.id));
 
   for (const instr of jsx) {
     const {value} = instr;
@@ -412,11 +427,30 @@ function emitUpdatedJsx(
       });
     }
 
+    let newChildren: Array<Place> | null = null;
+    if (value.children) {
+      newChildren = [];
+      for (const child of value.children) {
+        if (jsxIds.has(child.identifier.id)) {
+          newChildren.push({...child});
+          continue;
+        }
+
+        const newChild = oldToNewProps.get(child.identifier.id);
+        invariant(
+          newChild !== undefined,
+          `Expected a new prop for ${printIdentifier(child.identifier)}`,
+        );
+        newChildren.push({...newChild.place});
+      }
+    }
+
     newInstrs.push({
       ...instr,
       value: {
         ...value,
         props: newProps,
+        children: newChildren,
       },
     });
   }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.jsx-outlining-children.expect.md renamed to compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/jsx-outlining-with-non-jsx-children.expect.md

@@ -11,12 +11,14 @@ function Component({arr}) {
         return (
           <Bar key={id} x={x}>
             <Baz i={i}>Test</Baz>
+            <Foo k={i} />
           </Bar>
         );
       })}
     </>
   );
 }
+
 function Bar({x, children}) {
   return (
     <>
@@ -26,8 +28,17 @@ function Bar({x, children}) {
   );
 }
 
-function Baz({i}) {
-  return i;
+function Baz({i, children}) {
+  return (
+    <>
+      {i}
+      {children}
+    </>
+  );
+}
+
+function Foo({k}) {
+  return k;
 }
 
 function useX() {
@@ -53,11 +64,11 @@ function Component(t0) {
   if ($[0] !== arr || $[1] !== x) {
     let t2;
     if ($[3] !== x) {
-      t2 = (i, id) => (
-        <Bar key={id} x={x}>
-          <Baz i={i}>Test</Baz>
-        </Bar>
-      );
+      t2 = (i, id) => {
+        const t3 = "Test";
+        const T0 = _temp;
+        return <T0 i={i} t={t3} k={i} key={id} x={x} />;
+      };
       $[3] = x;
       $[4] = t2;
     } else {
@@ -80,6 +91,43 @@ function Component(t0) {
   }
   return t2;
 }
+function _temp(t0) {
+  const $ = _c(9);
+  const { i: i, t: t, k: k, x: x } = t0;
+  let t1;
+  if ($[0] !== i || $[1] !== t) {
+    t1 = <Baz i={i}>{t}</Baz>;
+    $[0] = i;
+    $[1] = t;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  let t2;
+  if ($[3] !== k) {
+    t2 = <Foo k={k} />;
+    $[3] = k;
+    $[4] = t2;
+  } else {
+    t2 = $[4];
+  }
+  let t3;
+  if ($[5] !== t1 || $[6] !== t2 || $[7] !== x) {
+    t3 = (
+      <Bar x={x}>
+        {t1}
+        {t2}
+      </Bar>
+    );
+    $[5] = t1;
+    $[6] = t2;
+    $[7] = x;
+    $[8] = t3;
+  } else {
+    t3 = $[8];
+  }
+  return t3;
+}
 
 function Bar(t0) {
   const $ = _c(3);
@@ -102,8 +150,28 @@ function Bar(t0) {
 }
 
 function Baz(t0) {
-  const { i } = t0;
-  return i;
+  const $ = _c(3);
+  const { i, children } = t0;
+  let t1;
+  if ($[0] !== children || $[1] !== i) {
+    t1 = (
+      <>
+        {i}
+        {children}
+      </>
+    );
+    $[0] = children;
+    $[1] = i;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+function Foo(t0) {
+  const { k } = t0;
+  return k;
 }
 
 function useX() {
@@ -118,4 +186,4 @@ export const FIXTURE_ENTRYPOINT = {
 ```
       
 ### Eval output
-(kind: ok) xfooxbar
+(kind: ok) xfooTestfooxbarTestbar

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.jsx-outlining-children.js renamed to compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/jsx-outlining-with-non-jsx-children.js

@@ -7,12 +7,14 @@ function Component({arr}) {
         return (
           <Bar key={id} x={x}>
             <Baz i={i}>Test</Baz>
+            <Foo k={i} />
           </Bar>
         );
       })}
     </>
   );
 }
+
 function Bar({x, children}) {
   return (
     <>
@@ -22,8 +24,17 @@ function Bar({x, children}) {
   );
 }
 
-function Baz({i}) {
-  return i;
+function Baz({i, children}) {
+  return (
+    <>
+      {i}
+      {children}
+    </>
+  );
+}
+
+function Foo({k}) {
+  return k;
 }
 
 function useX() {