fix hmr of mdx and outlineInfo; support used together with @vitejs/plugin-react-swc (#146)

Author: csr632

packages/playground/use-theme-doc/__tests__/hmr.spec.ts

@@ -77,6 +77,57 @@ test('hmr: edit file (md static data notation)', async ({
   ).toHaveCount(0)
 })
 
+test('hmr: edit md file content', async ({ page, fsUtils, testPlayground }) => {
+  // prepare locators first
+  const headingBeforeEdit = page
+    .locator('.markdown-body')
+    .getByRole('heading', { name: 'Heading one', exact: true })
+  const headingAfterEdit = page
+    .locator('.markdown-body')
+    .getByRole('heading', { name: 'Heading edited', exact: true })
+  // Also check the table-of-content
+  const outlineLinkBeforeEdit = page
+    .locator('.vp-local-outline')
+    .getByRole('link', { name: 'Heading one', exact: true })
+  const outlineLinkAfterEdit = page
+    .locator('.vp-local-outline')
+    .getByRole('link', { name: 'Heading edited', exact: true })
+  const counter = page.locator('.markdown-body').getByTestId('counter')
+  const counterStateText = counter.locator('span')
+  const counterButton = counter.getByRole('button', { name: 'add count' })
+
+  page.locator('.vp-local-sider >> text="Markdown Test Page1"').click()
+  await page.waitForURL('/md-test1')
+
+  // initial state
+  await expect(headingBeforeEdit).toHaveCount(1)
+  await expect(headingAfterEdit).toHaveCount(0)
+  await expect(outlineLinkBeforeEdit).toHaveCount(1)
+  await expect(outlineLinkAfterEdit).toHaveCount(0)
+  // update component state
+  await expect(counterStateText).toHaveText('Counter component: 0.')
+  await counterButton.click()
+  await expect(counterStateText).toHaveText('Counter component: 1.')
+
+  fsUtils.editFile('pages/md-test1$.md', (str) => {
+    return str.replace('# Heading one', '# Heading edited')
+  })
+
+  await expect(headingBeforeEdit).toHaveCount(0)
+  await expect(headingAfterEdit).toHaveCount(1)
+  await expect(outlineLinkBeforeEdit).toHaveCount(0)
+  await expect(outlineLinkAfterEdit).toHaveCount(1)
+  await expect(counterStateText).toHaveText('Counter component: 1.')
+
+  await testPlayground.restore()
+
+  await expect(headingBeforeEdit).toHaveCount(1)
+  await expect(headingAfterEdit).toHaveCount(0)
+  await expect(outlineLinkBeforeEdit).toHaveCount(1)
+  await expect(outlineLinkAfterEdit).toHaveCount(0)
+  await expect(counterStateText).toHaveText('Counter component: 1.')
+})
+
 test('hmr: delete file, add file', async ({ page, fsUtils }) => {
   const page2FileContent = fsUtils.readFile('pages/page2$.md')
   await expect(

packages/playground/use-theme-doc/package.json

@@ -23,6 +23,7 @@
   "devDependencies": {
     "@types/react": "^18.2.13",
     "@vitejs/plugin-react": "^4.0.1",
+    "@vitejs/plugin-react-swc": "^3.3.2",
     "rimraf": "^4.4.1",
     "serve": "^14.2.0",
     "vite": "^4.3.9",

packages/playground/use-theme-doc/pages/Counter.tsx

@@ -0,0 +1,15 @@
+import React, { useState } from 'react'
+
+interface Props {}
+
+const Counter: React.FC<Props> = (props) => {
+  const [count, setCount] = useState(0)
+  return (
+    <div data-testid="counter">
+      <span>Counter component: {count}.</span>
+      <button onClick={() => setCount((v) => v + 1)}>add count</button>
+    </div>
+  )
+}
+
+export default Counter

packages/playground/use-theme-doc/pages/md-test1$.md

@@ -2,6 +2,12 @@
 
 > markdown render test data from: https://github.com/fullpipe/markdown-test-page/blob/master/test-page.md
 
+# Counter
+
+import Counter from './Counter';
+
+<Counter />
+
 # Heading one
 
 Sint sit cillum pariatur eiusmod nulla pariatur ipsum. Sit laborum anim qui mollit tempor pariatur nisi minim dolor. Aliquip et adipisicing sit sit fugiat commodo id sunt. Nostrud enim ad commodo incididunt cupidatat in ullamco ullamco Lorem cupidatat velit enim et Lorem. Ut laborum cillum laboris fugiat culpa sint irure do reprehenderit culpa occaecat. Exercitation esse mollit tempor magna aliqua in occaecat aliquip veniam reprehenderit nisi dolor in laboris dolore velit.

packages/playground/use-theme-doc/vite.config.ts

@@ -1,6 +1,8 @@
 import { defineConfig } from 'vite'
 import * as path from 'path'
-import react from '@vitejs/plugin-react'
+// import react from '@vitejs/plugin-react'
+import react from "@vitejs/plugin-react-swc";
+
 import pages from 'vite-plugin-react-pages'
 
 export default defineConfig({

packages/react-pages/client.rollup.config.js

@@ -9,46 +9,60 @@ bundle client modules to reduce browser request during dev
 also prevent conflict npm packages with users (like jotai)
 */
 
-export default {
-  input: [
-    'src/client/entries/csr.tsx',
-    'src/client/entries/ssg-client.tsx',
-    'src/client/entries/ssg-server.tsx',
-  ],
-  output: {
-    dir: 'dist/client-bundles/entries',
-    entryFileNames: `[name].mjs`,
-    format: 'esm',
-    sourcemap: true,
+export default [
+  {
+    input: [
+      'src/client/entries/ssg-client.tsx',
+      'src/client/entries/ssg-server.tsx',
+    ],
+    ...config(),
   },
-  external: [],
-  plugins: [
-    {
-      name: 'client-external',
-      resolveId(source, importer) {
-        if (source.startsWith('/@react-pages/')) {
-          return {
-            id: source,
-            external: 'absolute',
+  // 1. stand alone bundle for dev, without code spliting. To reduce request waterfall
+  // 2. `hoistTransitiveImports` will add import (like `import '/@react-pages/pages'`) to the entry, which will mess up the hmr handling, so turn it off
+  {
+    input: ['src/client/entries/csr.tsx'],
+    ...config({ hoistTransitiveImports: false }),
+  },
+]
+
+function config({ hoistTransitiveImports } = {}) {
+  return {
+    output: {
+      dir: 'dist/client-bundles/entries',
+      entryFileNames: `[name].mjs`,
+      format: 'esm',
+      sourcemap: true,
+      hoistTransitiveImports,
+    },
+    external: [],
+    plugins: [
+      {
+        name: 'client-external',
+        resolveId(source, importer) {
+          if (source.startsWith('/@react-pages/')) {
+            return {
+              id: source,
+              external: 'absolute',
+            }
           }
-        }
+        },
       },
-    },
-    resolve({
-      // prevent bundling unexpected deps
-      resolveOnly: ['jotai'],
-      extensions,
-    }),
-    commonjs(),
-    babel({
-      babelHelpers: 'bundled',
-      extensions,
-      presets: [
-        '@babel/preset-typescript',
-        ['@babel/preset-react', { runtime: 'automatic' }],
-      ],
-      plugins: [],
-      configFile: false,
-    }),
-  ],
+      resolve({
+        // prevent bundling unexpected deps
+        resolveOnly: ['jotai'],
+        extensions,
+      }),
+      commonjs(),
+      babel({
+        babelHelpers: 'bundled',
+        extensions,
+        presets: [
+          '@babel/preset-typescript',
+          ['@babel/preset-react', { runtime: 'automatic' }],
+        ],
+        plugins: [],
+        configFile: false,
+      }),
+    ],
+  }
 }

packages/react-pages/src/client/state.ts

@@ -60,15 +60,15 @@ if (import.meta.hot) {
     setPages?.(module.default)
   })
 
-  let setAllPagesOutlines: SetAtom<any, void> | undefined
-  import.meta.hot!.accept('/@react-pages/allPagesOutlines', (module) => {
-    // console.log('@@hot update /@react-pages/allPagesOutlines', module)
-    if (!module) {
-      console.error('unexpected hot module', module)
-      return
-    }
-    setAllPagesOutlines?.(module)
-  })
+  // let setAllPagesOutlines: SetAtom<any, void> | undefined
+  // import.meta.hot!.accept('/@react-pages/allPagesOutlines', (module) => {
+  //   // console.log('@@hot update /@react-pages/allPagesOutlines', module)
+  //   if (!module) {
+  //     console.error('unexpected hot module', module)
+  //     return
+  //   }
+  //   setAllPagesOutlines?.(module)
+  // })
 
   const pagesAtom = atom(initialPages)
   const pagePathsAtom = atom(initialPagePaths.sort())
@@ -164,7 +164,7 @@ if (import.meta.hot) {
 
   useAllPagesOutlines = (timeout: number) => {
     const [data, set] = useAtom(allPagesOutlinesAtom)
-    setAllPagesOutlines = set
+    // setAllPagesOutlines = set
     useEffect(() => {
       setTimeout(() => {
         import('/@react-pages/allPagesOutlines').then((mod) => {

packages/react-pages/src/node/index.ts

@@ -1,6 +1,6 @@
 import * as path from 'path'
 import type { PluggableList } from 'unified'
-import type { Plugin, IndexHtmlTransformContext } from 'vite'
+import type { Plugin, IndexHtmlTransformContext, PluginOption } from 'vite'
 import type { OutputPlugin } from 'rollup'
 import type { staticSiteGenerationConfig } from './types'
 
@@ -269,21 +269,25 @@ function moveScriptTagToBodyEnd(
 
 export default async function setupPlugins(
   vpConfig: PluginConfig = {}
-): Promise<Plugin[]> {
+): Promise<PluginOption[]> {
   // use dynamic import so that it supports node commonjs
   const mdx = await import('@mdx-js/rollup')
+  const mdxPlugin = mdx.default({
+    remarkPlugins: await getRemarkPlugins(),
+    rehypePlugins: await getRehypePlugins(),
+    // treat .md as mdx
+    mdExtensions: [],
+    mdxExtensions: ['.md', '.mdx'],
+    providerImportSource: '@mdx-js/react',
+  })
   return [
-    mdx.default({
-      remarkPlugins: await getRemarkPlugins(),
-      rehypePlugins: await getRehypePlugins(),
-      // treat .md as mdx
-      mdExtensions: [],
-      mdxExtensions: ['.md', '.mdx'],
-      providerImportSource: '@mdx-js/react',
-    }),
+    {
+      ...mdxPlugin,
+      enforce: 'pre',
+    },
     createMdxTransformPlugin(),
     pluginFactory(vpConfig),
-  ] as Plugin[]
+  ]
 }
 
 function getRemarkPlugins(): Promise<PluggableList> {
@@ -314,34 +318,36 @@ function getRehypePlugins(): Promise<PluggableList> {
 function createMdxTransformPlugin(): Plugin {
   let vitePluginReactTrasnform: Plugin['transform'] | undefined
   return {
-    name: 'vite-pages:mdx-transform',
+    name: 'vite-pages:mdx-fast-refresh',
+    apply: 'serve',
     configResolved: ({ plugins }) => {
       // find this plugin to call it's transform function:
       // https://github.com/vitejs/vite-plugin-react/blob/b647e74c38565696bd6fb931b8bd9ac7f3bebe88/packages/plugin-react/src/index.ts#L206
+      // or https://github.com/vitejs/vite-plugin-react-swc/blob/95e991914322e7b011d1c8d18d501b9eee21adaa/src/index.ts#L111
       vitePluginReactTrasnform = plugins.find(
         (p) =>
-          p.name === 'vite:react-babel' && typeof p.transform === 'function'
+          (p.name === 'vite:react-babel' &&
+            typeof p.transform === 'function') ||
+          (p.name === 'vite:react-swc' && typeof p.transform === 'function')
       )?.transform
       if (!vitePluginReactTrasnform) {
         throw new Error(
-          `Can't find an instance of @vitejs/plugin-react. You should apply this plugin to make mdx work.`
+          `Can't find an instance of @vitejs/plugin-react or @vitejs/plugin-react-swc. You should apply either of these plugins to make mdx work.`
         )
       }
     },
     transform: (code, id, options) => {
-      const [filepath, querystring = ''] = id.split('?')
+      const [filepath, ...qs] = id.split('?')
       if (
         filepath.match(/\.mdx?$/) &&
         !id.startsWith(OUTLINE_INFO_MODULE_ID_PREFIX)
       ) {
-        // make @vitejs/plugin-react treat the "output of @mdx-js/rollup transform" like a jsx file
-        // https://github.com/vitejs/vite-plugin-react/blob/b647e74c38565696bd6fb931b8bd9ac7f3bebe88/packages/plugin-react/src/index.ts#L215
-        let newId
-        if (querystring) {
-          newId = id + '&ext=.jsx'
-        } else {
-          newId = id + '?ext=.jsx'
-        }
+        // turn file path like `/path/to/md-file$.md` into `/path/to/md-file$.jsx`
+        // make vite-plugin-react transform "the output of @mdx-js/rollup" like a jsx file
+        // https://github.com/vitejs/vite-plugin-react/blob/caa9b5330092c70288fcb94ceb96ca42438df2a2/packages/plugin-react/src/index.ts#L170
+        const newFilePath = filepath.replace(/\.mdx?$/, '.jsx')
+        const newId = [newFilePath, ...qs].join('?')
+        
         return (vitePluginReactTrasnform as any)(code, newId, options)
       }
     },