From ce86503a1a2b5da84b908b3a800cef18c8606acd Mon Sep 17 00:00:00 2001
From: Jihchi Lee <achi@987.tw>
Date: Sat, 23 Sep 2023 09:31:26 +0000
Subject: [PATCH] feat: add support for non-dev jsx runtime

---
 packages/plugin-react/src/index.ts         |  9 ++++++---
 playground/react/App.jsx                   |  2 ++
 playground/react/__tests__/react.spec.ts   | 23 ++++++++++++++++++++++
 playground/react/hmr/jsx-import-runtime.js |  8 ++++++++
 4 files changed, 39 insertions(+), 3 deletions(-)
 create mode 100644 playground/react/hmr/jsx-import-runtime.js

diff --git a/packages/plugin-react/src/index.ts b/packages/plugin-react/src/index.ts
index 4208c2e34..814674dfe 100644
--- a/packages/plugin-react/src/index.ts
+++ b/packages/plugin-react/src/index.ts
@@ -98,7 +98,9 @@ export default function viteReact(opts: Options = {}): PluginOption[] {
   // Provide default values for Rollup compat.
   let devBase = '/'
   const filter = createFilter(opts.include ?? defaultIncludeRE, opts.exclude)
-  const devRuntime = `${opts.jsxImportSource ?? 'react'}/jsx-dev-runtime`
+  const jsxImportSource = opts.jsxImportSource ?? 'react'
+  const jsxImportRuntime = `${jsxImportSource}/jsx-runtime`
+  const jsxImportDevRuntime = `${jsxImportSource}/jsx-dev-runtime`
   let isProduction = true
   let projectRoot = process.cwd()
   let skipFastRefresh = false
@@ -188,7 +190,8 @@ export default function viteReact(opts: Options = {}): PluginOption[] {
         (isJSX ||
           (opts.jsxRuntime === 'classic'
             ? importReactRE.test(code)
-            : code.includes(devRuntime)))
+            : code.includes(jsxImportDevRuntime) ||
+              code.includes(jsxImportRuntime)))
       if (useFastRefresh) {
         plugins.push([
           await loadPlugin('react-refresh/babel'),
@@ -265,7 +268,7 @@ export default function viteReact(opts: Options = {}): PluginOption[] {
         // We can't add `react-dom` because the dependency is `react-dom/client`
         // for React 18 while it's `react-dom` for React 17. We'd need to detect
         // what React version the user has installed.
-        include: ['react', devRuntime],
+        include: ['react', jsxImportDevRuntime, jsxImportRuntime],
       },
       resolve: {
         dedupe: ['react', 'react-dom'],
diff --git a/playground/react/App.jsx b/playground/react/App.jsx
index 83f4cc07e..95097029e 100644
--- a/playground/react/App.jsx
+++ b/playground/react/App.jsx
@@ -2,6 +2,7 @@ import { useState } from 'react'
 import Button from 'jsx-entry'
 import Dummy from './components/Dummy?qs-should-not-break-plugin-react'
 import Parent from './hmr/parent'
+import { JsxImportRuntime } from './hmr/jsx-import-runtime'
 import { CountProvider } from './context/CountProvider'
 import { ContextButton } from './context/ContextButton'
 
@@ -37,6 +38,7 @@ function App() {
 
       <Dummy />
       <Parent />
+      <JsxImportRuntime />
       <Button>button</Button>
     </div>
   )
diff --git a/playground/react/__tests__/react.spec.ts b/playground/react/__tests__/react.spec.ts
index e1e373ace..54180219f 100644
--- a/playground/react/__tests__/react.spec.ts
+++ b/playground/react/__tests__/react.spec.ts
@@ -115,4 +115,27 @@ if (!isBuild) {
       'context provider updated',
     )
   })
+
+  test('should hmr files with "react/jsx-runtime"', async () => {
+    expect(await page.textContent('#state-button')).toMatch('count is: 0')
+    await page.click('#state-button')
+    expect(await page.textContent('#state-button')).toMatch('count is: 1')
+
+    await untilBrowserLogAfter(
+      () =>
+        editFile('hmr/jsx-import-runtime.js', (code) =>
+          code.replace(
+            'JSX import runtime works',
+            'JSX import runtime updated',
+          ),
+        ),
+      ['[vite] hot updated: /hmr/jsx-import-runtime.js'],
+    )
+    await untilUpdated(
+      () => page.textContent('#jsx-import-runtime'),
+      'JSX import runtime updated',
+    )
+
+    expect(await page.textContent('#state-button')).toMatch('count is: 1')
+  })
 }
diff --git a/playground/react/hmr/jsx-import-runtime.js b/playground/react/hmr/jsx-import-runtime.js
new file mode 100644
index 000000000..27f1529dd
--- /dev/null
+++ b/playground/react/hmr/jsx-import-runtime.js
@@ -0,0 +1,8 @@
+import * as JsxRuntime from 'react/jsx-runtime'
+
+export function JsxImportRuntime() {
+  return JsxRuntime.jsx('p', {
+    id: 'jsx-import-runtime',
+    children: 'JSX import runtime works',
+  })
+}