[compiler] Preserve JSX pragmas before generated imports

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Imports.ts

@@ -302,9 +302,43 @@ export function addImportsToProgram(
       }
     }
   }
+  preserveLeadingJsxPragmaCommentsOnInsertedStatements(path.node, stmts);
   path.unshiftContainer('body', stmts);
 }
 
+function preserveLeadingJsxPragmaCommentsOnInsertedStatements(
+  program: t.Program,
+  stmts: Array<t.ImportDeclaration | t.VariableDeclaration>,
+): void {
+  const firstInsertedStmt = stmts[0];
+  const firstExistingStmt = program.body[0];
+  if (
+    firstInsertedStmt == null ||
+    firstExistingStmt == null ||
+    firstExistingStmt.leadingComments == null
+  ) {
+    return;
+  }
+  const remainingComments = [];
+  const jsxPragmaComments = [];
+  for (const comment of firstExistingStmt.leadingComments) {
+    if (/@jsx(?:ImportSource|Runtime|Frag)?\b/.test(comment.value)) {
+      jsxPragmaComments.push(comment);
+    } else {
+      remainingComments.push(comment);
+    }
+  }
+  if (jsxPragmaComments.length === 0) {
+    return;
+  }
+  firstInsertedStmt.leadingComments = [
+    ...(firstInsertedStmt.leadingComments ?? []),
+    ...jsxPragmaComments,
+  ];
+  firstExistingStmt.leadingComments =
+    remainingComments.length === 0 ? null : remainingComments;
+}
+
 /*
  * Matches `import { ... } from <moduleName>;`
  * but not `import * as React from <moduleName>;`

compiler/packages/babel-plugin-react-compiler/src/__tests__/Imports-test.ts

@@ -0,0 +1,30 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import invariant from 'invariant';
+import {runBabelPluginReactCompiler} from '../Babel/RunReactCompilerBabelPlugin';
+
+it('preserves JSX pragmas before generated runtime imports', () => {
+  const result = runBabelPluginReactCompiler(
+    `/* @jsxImportSource custom-jsx */
+
+import {useState} from 'react';
+
+export default function Component({children}) {
+  const [count] = useState(0);
+  return <div data-count={count}>{children}</div>;
+}`,
+    'test.tsx',
+    'typescript',
+    {panicThreshold: 'all_errors'},
+  );
+
+  invariant(result.code != null, 'Expected Babel transform to emit code');
+  expect(result.code.indexOf('@jsxImportSource custom-jsx')).toBeLessThan(
+    result.code.indexOf('react/compiler-runtime'),
+  );
+});