From cdde15efe189e8bebe227b7555e7cc95ad74deab Mon Sep 17 00:00:00 2001
From: Joseph Savona <6425824+josephsavona@users.noreply.github.com>
Date: Fri, 18 Oct 2024 11:27:48 -0700
Subject: [PATCH] [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.
---
 .../src/Optimization/InlineJsxTransform.ts    | 11 +++
 .../compiler/inline-jsx-transform.expect.md   | 96 +++++++++++++------
 .../fixtures/compiler/inline-jsx-transform.js | 10 +-
 3 files changed, 83 insertions(+), 34 deletions(-)

diff --git a/compiler/packages/babel-plugin-react-compiler/src/Optimization/InlineJsxTransform.ts b/compiler/packages/babel-plugin-react-compiler/src/Optimization/InlineJsxTransform.ts
index 396e6ad1be..89efa78469 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Optimization/InlineJsxTransform.ts
+++ b/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);
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.expect.md
index 6dd899d5c7..2078575e83 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.expect.md
+++ b/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>
\ No newline at end of file
+(kind: ok) <div><span class="abc">Hello world</span><div>abc</div></div>
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.js
index 4a9d53b6f4..6fe9553dcd 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.js
+++ b/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" />
     </>
   );
 }