experiment with ssr hmr

src/App.tsx

@@ -1,4 +1,5 @@
 import { useState } from 'react'
+import { AppDep1, AppDep2 } from './AppDep'
 
 export function App() {
   const [count, setCount] = useState(0)
@@ -11,6 +12,10 @@ export function App() {
           count is {count}
         </button>
       </div>
+      <div>
+        <AppDep1 />
+        <AppDep2 />
+      </div>
     </>
   )
 }

src/AppDep.tsx

@@ -0,0 +1,9 @@
+import { AppDepDep } from "./AppDepDep"
+
+export function AppDep1() {
+  return <div>AppDep1</div>
+}
+
+export function AppDep2() {
+  return <div>AppDep2 - <AppDepDep /></div>
+}

src/AppDepDep.tsx

@@ -0,0 +1,3 @@
+export function AppDepDep() {
+  return <span>AppDepDep</span>;
+}

src/hmr.ts

@@ -0,0 +1,73 @@
+// runtime part of `ssrHmrPlugin` in `vite.config.ts`
+
+// simple automatic HMR for SSR
+// - full reload if export names changed
+// - otherwise reassign all exports
+
+// https://github.com/vitejs/vite/pull/12165
+// https://github.com/hi-ogawa/js-utils/blob/main/packages/tiny-refresh/src/runtime.ts
+
+const REGISTRY_KEY = Symbol.for("simple-hmr");
+
+export interface ViteHot {
+  data: {
+    [REGISTRY_KEY]?: Registry;
+  };
+  accept: (onNewModule: (exports?: unknown) => void) => void;
+  invalidate: (message?: string) => void;
+}
+
+interface Export {
+  value: unknown;
+  update: (next: unknown) => void;
+}
+
+interface Registry {
+  exports: Record<string, Export>;
+  // keep track of all exports of hot update history
+  // since currently "writer" is responsible to keep old module up-to-date
+  // while each export could be used by other module at any point in time
+  // but this approach obviously leaks memory indefinitely
+  // (alternative is to let "reader" be responsible for looking up latest module using proxy)
+  history: Record<string, Export>[];
+}
+
+export function createRegistry(): Registry {
+  const exports = {};
+  return { exports, history: [exports] };
+}
+
+function patchRegistry(current: Registry, next: Registry): boolean {
+  // replace all exports in history or full reload
+  const keys = [
+    ...new Set([...Object.keys(current.exports), ...Object.keys(next.exports)]),
+  ];
+  const mismatches = keys.filter(
+    (key) => !(key in current.exports && key in next.exports)
+  );
+  if (mismatches.length > 0) {
+    console.log("[simple-hmr] mismatch: ", mismatches.join(", "));
+    return false;
+  }
+  for (const key of keys) {
+    console.log("[simple-hmr]", key);
+    for (const e of current.history) {
+      e[key].update(next.exports[key].value);
+    }
+  }
+  next.history = current.history;
+  next.history.push(next.exports);
+  return true;
+}
+
+export function setupHot(hot: ViteHot, registry: Registry) {
+  hot.data[REGISTRY_KEY] = registry;
+
+  hot.accept((newExports) => {
+    const current = hot.data[REGISTRY_KEY];
+    const ok = newExports && current && patchRegistry(registry, current);
+    if (!ok) {
+      hot.invalidate();
+    }
+  });
+}

vite.config.ts

@@ -1,25 +1,34 @@
-import { defineConfig, type Plugin, type Connect } from 'vite'
-import react from '@vitejs/plugin-react'
+import {
+  defineConfig,
+  type Plugin,
+  type Connect,
+  FilterPattern,
+  createFilter,
+} from "vite";
+import react from "@vitejs/plugin-react";
 
 // https://vitejs.dev/config/
 export default defineConfig({
   clearScreen: false,
   plugins: [
     react(),
+    ssrHmrPlugin({
+      include: ["**/src/**/*.tsx"],
+    }),
     devPlugin({
       entry: "/src/server.ts",
       useViteRuntime: true,
-    })
-  ]
-})
+    }),
+  ],
+});
 
 // vavite style dev server with ViteRuntime.executeEntrypoint
 // https://github.com/hi-ogawa/vite-plugins/pull/156
 function devPlugin({
   entry,
   useViteRuntime,
 }: {
-  entry: string
+  entry: string;
   useViteRuntime?: boolean;
 }): Plugin {
   return {
@@ -48,6 +57,89 @@ function devPlugin({
         }
       };
       return () => server.middlewares.use(handler);
+    },
+  };
+}
+
+function ssrHmrPlugin(pluginOpts: {
+  include?: FilterPattern;
+  exclude?: FilterPattern;
+}): Plugin {
+  const filter = createFilter(
+    pluginOpts.include,
+    pluginOpts.exclude ?? ["**/node_modules/**"]
+  );
+
+  return {
+    name: ssrHmrPlugin.name,
+
+    apply: "serve",
+
+    transform(code, id, options) {
+      if (options.ssr && filter(id)) {
+        return ssrHmrTransform(code, id);
+      }
+      return;
+    },
+  };
+}
+
+function ssrHmrTransform(code: string, id: string): string {
+  // transform to inject something like below
+  /*
+    if (import.meta.env.SSR && import.meta.hot) {
+      const $$hmr = await import("./hmr");
+      const $$registry = $$hmr.createRegistry();
+
+      $$registry.exports["App"] = {
+        value: App,
+        update: ($$next) => {
+          // @ts-ignore
+          App = $$next;
+        }
+      };
+
+      $$hmr.setupHot(import.meta.hot, $$registry);
     }
+  */
+
+  // TODO: use vite/rollup parser
+  // TODO: replace `export const` with `export let` for reassignment
+  // TODO: magic-string + sourcemap
+
+  // extract named exports
+  const matches = code.matchAll(/export\s+(function|let)\s+(\w+)\b/g);
+  const exportNames = Array.from(matches).map((m) => m[2]);
+  if (0) {
+    console.log({ id }, exportNames)
+  }
+
+  if (exportNames.length === 0) {
+    return undefined;
   }
+
+  // append runtime in footer
+  const parts = exportNames.map(
+    (name) => `
+  $$registry.exports["${name}"] = {
+    value: ${name},
+    update: ($$next) => {
+      ${name} = $$next;
+    }
+  };
+`
+  );
+
+  const footer = `
+if (import.meta.env.SSR && import.meta.hot) {
+  const $$hmr = await import("/src/hmr");
+  const $$registry = $$hmr.createRegistry();
+
+${parts.join("\n")}
+
+  $$hmr.setupHot(import.meta.hot, $$registry);
+  import.meta.hot.accept; // dummy code for vite's "hot.accept" detection
+}
+`;
+  return code + footer;
 }