[compiler] InlineJSXTransform transforms jsx inside function expressions (#31282)

InlineJSXTransform wasn't traversing into function expressions or object
methods, so any JSX inside such functions wouldn't have gotten inlined.
This PR updates to traverse nested functions to transform all JSX within
a hook or component.

Note that this still doesn't transform JSX outside of components or
hooks, ie in standalone render helpers.

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

@@ -405,6 +405,17 @@ export function inlineJsxTransform(
           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);

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

@@ -26,11 +26,15 @@ function ParentAndRefAndKey(props) {
 }
 
 function ParentAndChildren(props) {
+  const render = () => {
+    return <div key="d">{props.foo}</div>;
+  };
   return (
     <Parent>
       <Child key="a" {...props} />
       <Child key="b">
-        <GrandChild className={props.foo} {...props} />
+        <GrandChild key="c" className={props.foo} {...props} />
+        {render()}
       </Child>
     </Parent>
   );
@@ -40,8 +44,8 @@ const propsToSpread = {a: 'a', b: 'b', c: 'c'};
 function PropsSpread() {
   return (
     <>
-      <Test {...propsToSpread} />
-      <Test {...propsToSpread} a="z" />
+      <Test key="a" {...propsToSpread} />
+      <Test key="b" {...propsToSpread} a="z" />
     </>
   );
 }
@@ -151,37 +155,30 @@ function ParentAndRefAndKey(props) {
 }
 
 function ParentAndChildren(props) {
-  const $ = _c2(7);
+  const $ = _c2(14);
   let t0;
-  if ($[0] !== props) {
-    t0 = {
+  if ($[0] !== props.foo) {
+    t0 = () => ({
       $$typeof: Symbol.for("react.transitional.element"),
-      type: Child,
+      type: "div",
       ref: null,
-      key: "a",
-      props: props,
-    };
-    $[0] = props;
+      key: "d",
+      props: { children: props.foo },
+    });
+    $[0] = props.foo;
     $[1] = t0;
   } else {
     t0 = $[1];
   }
+  const render = t0;
   let t1;
   if ($[2] !== props) {
     t1 = {
       $$typeof: Symbol.for("react.transitional.element"),
       type: Child,
       ref: null,
-      key: "b",
-      props: {
-        children: {
-          $$typeof: Symbol.for("react.transitional.element"),
-          type: GrandChild,
-          ref: null,
-          key: null,
-          props: { className: props.foo, ...props },
-        },
-      },
+      key: "a",
+      props: props,
     };
     $[2] = props;
     $[3] = t1;
@@ -189,21 +186,58 @@ function ParentAndChildren(props) {
     t1 = $[3];
   }
   let t2;
-  if ($[4] !== t0 || $[5] !== t1) {
+  if ($[4] !== props) {
     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) {
+    t3 = render();
+    $[6] = render;
+    $[7] = t3;
+  } else {
+    t3 = $[7];
+  }
+  let t4;
+  if ($[8] !== t2 || $[9] !== t3) {
+    t4 = {
+      $$typeof: Symbol.for("react.transitional.element"),
+      type: Child,
+      ref: null,
+      key: "b",
+      props: { children: [t2, t3] },
+    };
+    $[8] = t2;
+    $[9] = t3;
+    $[10] = t4;
+  } else {
+    t4 = $[10];
+  }
+  let t5;
+  if ($[11] !== t1 || $[12] !== t4) {
+    t5 = {
       $$typeof: Symbol.for("react.transitional.element"),
       type: Parent,
       ref: null,
       key: null,
-      props: { children: [t0, t1] },
+      props: { children: [t1, t4] },
     };
-    $[4] = t0;
-    $[5] = t1;
-    $[6] = t2;
+    $[11] = t1;
+    $[12] = t4;
+    $[13] = t5;
   } else {
-    t2 = $[6];
+    t5 = $[13];
   }
-  return t2;
+  return t5;
 }
 
 const propsToSpread = { a: "a", b: "b", c: "c" };
@@ -222,14 +256,14 @@ function PropsSpread() {
             $$typeof: Symbol.for("react.transitional.element"),
             type: Test,
             ref: null,
-            key: null,
+            key: "a",
             props: propsToSpread,
           },
           {
             $$typeof: Symbol.for("react.transitional.element"),
             type: Test,
             ref: null,
-            key: null,
+            key: "b",
             props: { ...propsToSpread, a: "z" },
           },
         ],
@@ -250,4 +284,4 @@ export const FIXTURE_ENTRYPOINT = {
 ```
       
 ### Eval output
-(kind: ok) <div><span class="abc">Hello world</span></div>
+(kind: ok) <div><span class="abc">Hello world</span><div>abc</div></div>

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

@@ -22,11 +22,15 @@ function ParentAndRefAndKey(props) {
 }
 
 function ParentAndChildren(props) {
+  const render = () => {
+    return <div key="d">{props.foo}</div>;
+  };
   return (
     <Parent>
       <Child key="a" {...props} />
       <Child key="b">
-        <GrandChild className={props.foo} {...props} />
+        <GrandChild key="c" className={props.foo} {...props} />
+        {render()}
       </Child>
     </Parent>
   );
@@ -36,8 +40,8 @@ const propsToSpread = {a: 'a', b: 'b', c: 'c'};
 function PropsSpread() {
   return (
     <>
-      <Test {...propsToSpread} />
-      <Test {...propsToSpread} a="z" />
+      <Test key="a" {...propsToSpread} />
+      <Test key="b" {...propsToSpread} a="z" />
     </>
   );
 }