feat: efficient HMR updates

This PR introduces new React hooks for subscribing to hot-reloadable data provided by the "vite-plugin-react-pages/dist/client" module. These hooks include:
  - useStaticData: Subscribe to the `.staticData` of a specific page or all pages.
  - usePagePaths: Subscribe to the array of paths for all pages.
  - usePageModule: Subscribe to the module referred to by the `dataPath` of a specific page.
  - useTheme: Subscribe to the `Theme` component.

In most cases, `useStaticData` is the only hook used by the application/theme. It replaces the `staticData` prop previously passed to the `Theme` component. The idea is to depend on static data where it's needed, instead of having to re-render the entire application when the static data is updated.

By calling `useStaticData()`, the component subscribes to **all static data.** By calling `useStaticData(path)`, the component only re-renders when the static data of that specific path is updated. You can even pass a selector if you only need a specific piece of static data. In the example below, the caller only re-renders when the page title is changed.

```ts
const title = useStaticData(path, data => data.main.title)
```

The other hooks are mostly for the internal infrastructure of react-pages.

`useTheme` is used by the `PageLoader` component (which is now wrapped with `React.memo` to avoid re-rendering when a page is added or removed).

`usePagePaths` is used by the `App` component for updating the page routes when a page is added or removed.

`usePageModule` is used by the `useAppState` hook (which I've rewritten) for reloading a page when its `dataPath` is changed. It even uses the `/404` page if a path does not exist. Once the path *does* exist, HMR will re-render the `PageLoader` automatically with a new `usePageModule` promise.

How does this all work?
---

When HMR is enabled, we use a library called [jotai](https://npmjs.org/jotai), which lets us re-render components only when the data they depend on is changed. The `client/state.ts` module (with access to our Jotai atoms) plugs into Vite's HMR API. When the `@!virtual-modules/pages` module is updated, we diff it with its previous module to determine what has changed (see `setPagesAtom` in `client/state.ts`).

In production, all that logic is tree-shaked, and we're left with simple hooks that return immutable data.

packages/react-pages/client.d.ts

@@ -4,7 +4,6 @@ import React from 'react'
 export type Theme = React.ComponentType<ThemeProps>
 
 export interface ThemeProps {
-  readonly staticData: PagesStaticData
   readonly loadedData: PagesLoaded
   readonly loadState: LoadState
 }

packages/react-pages/package.json

@@ -49,13 +49,15 @@
     "@babel/preset-typescript": "^7.12.17",
     "@rollup/plugin-babel": "^5.3.0",
     "chalk": "^4.1.0",
+    "dequal": "^2.0.2",
     "enhanced-resolve": "^5.7.0",
     "find-up": "^5.0.0",
     "fs-extra": "^9.1.0",
     "glob": "^7.1.6",
     "globby": "^11.0.2",
     "gray-matter": "^4.0.2",
     "jest-docblock": "^26.0.0",
+    "jotai": "^0.15.1",
     "mini-debounce": "^1.0.8",
     "minimist": "^1.2.5",
     "ora": "^5.3.0",

packages/react-pages/src/client/App.tsx

@@ -1,14 +1,18 @@
 import React from 'react'
 import { Switch, Route } from 'react-router-dom'
+import { usePagePaths } from './state'
 import PageLoader from './PageLoader'
 
-import pages from '@!virtual-modules/pages'
-import Theme from '@!virtual-modules/theme'
-
 const App = () => {
-  const pageRoutes = Object.keys(pages)
+  const pageRoutes = usePagePaths()
     .filter((path) => path !== '/404')
-    .map((path) => getPageRoute(path, pages[path].staticData))
+    .map((path) => (
+      // avoid re-mount layout component
+      // https://github.com/ReactTraining/react-router/issues/3928#issuecomment-284152397
+      <Route key="same" exact path={path}>
+        <PageLoader routePath={path} />
+      </Route>
+    ))
 
   return (
     <Switch>
@@ -18,32 +22,11 @@ const App = () => {
         path="*"
         render={({ match }) => {
           // 404
-          return (
-            <PageLoader Theme={Theme} pages={pages} routePath={match.url} />
-          )
+          return <PageLoader routePath={match.url} />
         }}
       />
     </Switch>
   )
 }
 
 export default App
-
-function getPageRoute(path: string, staticData: any) {
-  if (!pages[path]) {
-    throw new Error(`page not exist. route path: ${path}`)
-  }
-  return (
-    <Route
-      // avoid re-mount layout component
-      // https://github.com/ReactTraining/react-router/issues/3928#issuecomment-284152397
-      key="same"
-      exact
-      path={path}
-      // not used for now
-      {...staticData._routeConfig}
-    >
-      <PageLoader Theme={Theme} pages={pages} routePath={path} />
-    </Route>
-  )
-}

packages/react-pages/src/client/PageLoader.tsx

@@ -1,34 +1,18 @@
-import React, { useContext, useMemo } from 'react'
-import type { PagesStaticData, PagesInternal, Theme } from '../../client'
+import React, { useContext } from 'react'
 import { dataCacheCtx } from './ssr/ctx'
+import { useTheme } from './state'
 import useAppState from './useAppState'
 
 interface Props {
-  readonly Theme: Theme
-  readonly pages: PagesInternal
-  readonly routePath: string
+  routePath: string
 }
 
-const PageLoader = ({ pages, routePath: routePathFromProps, Theme }: Props) => {
+const PageLoader = React.memo(({ routePath }: Props) => {
+  const Theme = useTheme()
+  const loadState = useAppState(routePath)
   const dataCache = useContext(dataCacheCtx)
-  const loadState = useAppState(pages, routePathFromProps)
 
-  const pagesStaticData = useMemo(() => getPublicPages(pages), [pages])
-
-  return (
-    <Theme
-      loadState={loadState}
-      loadedData={dataCache}
-      staticData={pagesStaticData}
-    />
-  )
-}
+  return <Theme loadState={loadState} loadedData={dataCache} />
+})
 
 export default PageLoader
-
-// filter out internal fields inside pages
-function getPublicPages(pages: PagesInternal): PagesStaticData {
-  return Object.fromEntries(
-    Object.entries(pages).map(([path, { staticData }]) => [path, staticData])
-  )
-}

packages/react-pages/src/client/index.tsx

@@ -1 +1,2 @@
 export type { Theme } from '../../client'
+export { useStaticData } from './state'

packages/react-pages/src/client/main.tsx

@@ -1,11 +1,15 @@
 import React from 'react'
 import ReactDOM from 'react-dom'
+import { Provider as Jotai } from 'jotai'
 import SSRContextProvider from './SSRContextProvider'
 import App from './App'
 
+let app = <App />
+if (import.meta.hot) {
+  app = <Jotai>{app}</Jotai>
+}
+
 ReactDOM.render(
-  <SSRContextProvider>
-    <App />
-  </SSRContextProvider>,
+  <SSRContextProvider>{app}</SSRContextProvider>,
   document.getElementById('root')
 )

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

@@ -0,0 +1,155 @@
+import { useMemo } from 'react'
+import { dequal } from 'dequal'
+import type { SetAtom } from 'jotai/core/types'
+import { atom, useAtom } from 'jotai'
+import { atomFamily, useAtomValue, useUpdateAtom } from 'jotai/utils'
+import type { PageLoaded, PagesStaticData, Theme } from '../../client'
+
+export let useTheme: () => Theme
+export let usePagePaths: () => string[]
+export let usePageModule: (path: string) => Promise<PageModule> | undefined
+export let useStaticData: UseStaticData
+
+interface PageModule {
+  ['default']: PageLoaded
+}
+
+interface UseStaticData {
+  (): PagesStaticData
+  (path: string): Record<string, any>
+  <T>(path: string, selector: (staticData: Record<string, any>) => T): T
+}
+
+import initialPages from '@!virtual-modules/pages'
+import initialTheme from '@!virtual-modules/theme'
+
+const initialPagePaths = Object.keys(initialPages)
+
+// This HMR code assumes that our Jotai atoms are always managed
+// by the same Provider. It also mutates during render, which is
+// generally discouraged, but in this case it's okay.
+if (import.meta.hot) {
+  let setTheme: SetAtom<{ Theme: Theme }>
+  let setPages: SetAtom<PagesStaticData>
+
+  // Without /@id/ prefix, Vite will resolve them relative to __dirname.
+  import.meta.hot!.accept('/@id/@!virtual-modules/theme', (module) => {
+    setTheme(module.default)
+  })
+  import.meta.hot!.accept('/@id/@!virtual-modules/pages', (module) => {
+    setPages(module.default)
+  })
+
+  const themeAtom = atom({ Theme: initialTheme })
+  useTheme = () => {
+    const [{ Theme }, set] = useAtom(themeAtom)
+    setTheme = set
+    return Theme
+  }
+
+  const pagesAtom = atom(initialPages)
+  const pagePathsAtom = atom(initialPagePaths.sort())
+  const staticDataAtom = atom(toStaticData(initialPages))
+
+  const setPagesAtom = atom(null, (get, set, newPages: any) => {
+    let newStaticData: Record<string, any> | undefined
+
+    const pages = get(pagesAtom)
+    for (const path in newPages) {
+      const page = pages[path]
+      const newPage = newPages[path]
+
+      // Avoid changing the identity of `page.staticData` unless
+      // a change is detected. This prevents unnecessary renders
+      // of components that depend on `useStaticData(path)` call.
+      if (page && dequal(page.staticData, newPage.staticData)) {
+        newPage.staticData = page.staticData
+      } else {
+        newStaticData ??= {}
+        newStaticData[path] = newPage.staticData
+      }
+    }
+
+    // Update the `pagesAtom` every time, since no hook uses it directly.
+    set(pagesAtom, newPages)
+
+    // Avoid re-rendering `useStaticData()` callers if no data changed.
+    if (newStaticData) {
+      set(staticDataAtom, {
+        ...get(staticDataAtom),
+        ...newStaticData,
+      })
+    }
+
+    // Avoid re-rendering `usePagePaths()` callers if no paths were added/deleted.
+    const newPagePaths = Object.keys(newPages).sort()
+    if (!dequal(get(pagePathsAtom), newPagePaths)) {
+      set(pagePathsAtom, newPagePaths)
+    }
+  })
+
+  const dataPathAtoms = atomFamily((path: string) => (get) => {
+    const pages = get(pagesAtom)
+    const page = pages[path] || pages['/404']
+    return page?.dataPath || null
+  })
+
+  const emptyData: any = {}
+  const staticDataAtoms = atomFamily((path: string) => (get) => {
+    const pages = get(pagesAtom)
+    const page = pages[path] || pages['/404']
+    return page?.staticData || emptyData
+  })
+
+  usePagePaths = () => {
+    setPages = useUpdateAtom(setPagesAtom)
+    return useAtomValue(pagePathsAtom)
+  }
+
+  // This hook uses dynamic import with a variable, which is not supported
+  // by Rollup, but that's okay since HMR is for development only.
+  usePageModule = (pagePath) => {
+    const dataPath = useAtomValue(dataPathAtoms(pagePath))
+    return useMemo(() => {
+      return dataPath ? import(dataPath /* @vite-ignore */) : void 0
+    }, [dataPath])
+  }
+
+  useStaticData = (pagePath?: string, selector?: Function) => {
+    const staticData = pagePath ? staticDataAtoms(pagePath) : staticDataAtom
+    if (selector) {
+      const selection = useMemo(
+        () => atom((get) => selector(get(staticData))),
+        [staticData]
+      )
+      return useAtomValue(selection)
+    }
+    return useAtomValue(staticData)
+  }
+}
+
+// Static mode
+else {
+  useTheme = () => initialTheme
+  usePagePaths = () => initialPagePaths
+  usePageModule = (path) => {
+    const page = initialPages[path] || initialPages['/404']
+    return useMemo(() => page?.data(), [page])
+  }
+  useStaticData = (path?: string, selector?: Function) => {
+    if (path) {
+      const page = initialPages[path] || initialPages['/404']
+      const staticData = page?.staticData || {}
+      return selector ? selector(staticData) : staticData
+    }
+    return toStaticData(initialPages)
+  }
+}
+
+function toStaticData(pages: Record<string, any>) {
+  const staticData: Record<string, any> = {}
+  for (const path in pages) {
+    staticData[path] = pages[path].staticData
+  }
+  return staticData
+}