From 4beb1fd8bacacd3f4faeb180bca00ea0edcf27e4 Mon Sep 17 00:00:00 2001
From: Jack Pope <jackpope1@gmail.com>
Date: Mon, 18 Nov 2024 10:51:16 -0500
Subject: [PATCH] [compiler] Support enableRefAsProp in jsx transform (#31558)

Since `enableRefAsProp` shipped everywhere, the ReactElement
implementation on prod puts refs on both `element.ref` and
`element.props.ref`. Here we let the `ref` case fall through so its now
available on props, matching the JSX runtime.
---
 .../src/Optimization/InlineJsxTransform.ts    | 61 ++++++++++++-------
 .../compiler/inline-jsx-transform.expect.md   |  4 +-
 2 files changed, 41 insertions(+), 24 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 d97a4da1eccc4..2851ed7ee39a1 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Optimization/InlineJsxTransform.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Optimization/InlineJsxTransform.ts
@@ -559,28 +559,45 @@ function createPropsProperties(
   propAttributes.forEach(prop => {
     switch (prop.kind) {
       case 'JsxAttribute': {
-        if (prop.name === 'ref') {
-          refProperty = {
-            kind: 'ObjectProperty',
-            key: {name: 'ref', kind: 'string'},
-            type: 'property',
-            place: {...prop.place},
-          };
-        } else if (prop.name === 'key') {
-          keyProperty = {
-            kind: 'ObjectProperty',
-            key: {name: 'key', kind: 'string'},
-            type: 'property',
-            place: {...prop.place},
-          };
-        } else {
-          const attributeProperty: ObjectProperty = {
-            kind: 'ObjectProperty',
-            key: {name: prop.name, kind: 'string'},
-            type: 'property',
-            place: {...prop.place},
-          };
-          props.push(attributeProperty);
+        switch (prop.name) {
+          case 'key': {
+            keyProperty = {
+              kind: 'ObjectProperty',
+              key: {name: 'key', kind: 'string'},
+              type: 'property',
+              place: {...prop.place},
+            };
+            break;
+          }
+          case 'ref': {
+            /**
+             * In the current JSX implementation, ref is both
+             * a property on the element and a property on props.
+             */
+            refProperty = {
+              kind: 'ObjectProperty',
+              key: {name: 'ref', kind: 'string'},
+              type: 'property',
+              place: {...prop.place},
+            };
+            const refPropProperty: ObjectProperty = {
+              kind: 'ObjectProperty',
+              key: {name: 'ref', kind: 'string'},
+              type: 'property',
+              place: {...prop.place},
+            };
+            props.push(refPropProperty);
+            break;
+          }
+          default: {
+            const attributeProperty: ObjectProperty = {
+              kind: 'ObjectProperty',
+              key: {name: prop.name, kind: 'string'},
+              type: 'property',
+              place: {...prop.place},
+            };
+            props.push(attributeProperty);
+          }
         }
         break;
       }
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 f622b3a6fd667..91bd0ad0b750e 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
@@ -91,7 +91,7 @@ function Parent(t0) {
         type: "div",
         ref: ref,
         key: null,
-        props: { children: children },
+        props: { ref: ref, children: children },
       };
     }
     $[0] = children;
@@ -180,7 +180,7 @@ function ParentAndRefAndKey(props) {
         type: Parent,
         ref: testRef,
         key: "testKey",
-        props: { a: "a", b: { b: "b" }, c: C },
+        props: { a: "a", b: { b: "b" }, c: C, ref: testRef },
       };
     }
     $[0] = t0;