refactor(tiny-refresh): unify code for vite and webpack

packages/tiny-react/src/hmr/index.test.tsx

@@ -1,5 +1,5 @@
 import { describe, expect, it, vi } from "vitest";
-import { type ViteHot, createManager } from ".";
+import { type ViteHot, initialize } from ".";
 import { useEffect, useReducer } from "../hooks";
 import { render } from "../reconciler";
 import { sleepFrame } from "../test-utils";
@@ -45,7 +45,7 @@ describe("hmr", () => {
         return <div>1</div>;
       }
 
-      const manager = createManager(hot, runtime, false);
+      const manager = initialize(hot, runtime, { mode: "vite", debug: false });
       ChildExport = manager.wrap("Child", Child, "useEffect");
       manager.setup();
     }
@@ -78,7 +78,7 @@ describe("hmr", () => {
         return <div>2</div>;
       }
 
-      const manager = createManager(hot, runtime, false);
+      const manager = initialize(hot, runtime, { mode: "vite", debug: false });
       manager.wrap("Child", Child, "");
       manager.setup();
     }
@@ -119,7 +119,7 @@ describe("hmr", () => {
         return <div>3</div>;
       }
 
-      const manager = createManager(hot, runtime, false);
+      const manager = initialize(hot, runtime, { mode: "vite", debug: false });
       manager.wrap("Child", Child, "useEffect");
       manager.setup();
     }

packages/tiny-refresh/src/runtime.test.tsx

@@ -3,7 +3,7 @@
 import { act, cleanup, render } from "@testing-library/react";
 import React from "react";
 import { afterEach, describe, expect, it, vi } from "vitest";
-import { type ViteHot, createManager } from "./runtime";
+import { type ViteHot, initialize } from "./runtime";
 
 afterEach(cleanup);
 
@@ -37,7 +37,7 @@ describe("hmr", () => {
         return <div>1</div>;
       }
 
-      const manager = createManager(hot, React, false);
+      const manager = initialize(hot, React, { mode: "vite", debug: false });
       ChildExport = manager.wrap("Child", Child, "useEffect");
       manager.setup();
     }
@@ -70,7 +70,7 @@ describe("hmr", () => {
       function Child() {
         return <div>2</div>;
       }
-      const manager = createManager(hot, React, false);
+      const manager = initialize(hot, React, { mode: "vite", debug: false });
       manager.wrap("Child", Child, "");
       manager.setup();
     }
@@ -111,7 +111,7 @@ describe("hmr", () => {
         return <div>3</div>;
       }
 
-      const manager = createManager(hot, React, false);
+      const manager = initialize(hot, React, { mode: "vite", debug: false });
       manager.wrap("Child", Child, "useEffect");
       manager.setup();
     }
@@ -154,7 +154,7 @@ describe("hmr", () => {
         return <div>4</div>;
       }
 
-      const manager = createManager(hot, React, false);
+      const manager = initialize(hot, React, { mode: "vite", debug: false });
       manager.wrap("Child", Child, "useEffect");
       manager.setup();
     }
@@ -191,7 +191,7 @@ describe("hmr", () => {
         return <div>5</div>;
       }
 
-      const manager = createManager(hot, React, false);
+      const manager = initialize(hot, React, { mode: "vite", debug: false });
       manager.wrap("Child", Child, "");
       manager.setup();
     }

packages/tiny-refresh/src/runtime.ts

@@ -1,13 +1,24 @@
+import type { RefreshRuntimeOptions } from "./transform";
+
 const MANAGER_KEY = Symbol.for("tiny-refresh.manager");
 
 export interface ViteHot {
   accept: (onNewModule: (newModule?: unknown) => void) => void;
   invalidate: (message?: string) => void;
-  data: {
-    [MANAGER_KEY]?: Manager;
-  };
+  data: HotData;
+}
+
+interface WebpackHot {
+  accept: (cb?: () => void) => void;
+  invalidate: () => void;
+  dispose: (cb: (data: HotData) => void) => void;
+  data?: HotData;
 }
 
+type HotData = {
+  [MANAGER_KEY]?: Manager;
+};
+
 type FC = (props: any) => unknown;
 
 interface Runtime {
@@ -27,16 +38,14 @@ interface ComponentEntry {
 }
 
 // singleton per file
-export class Manager {
+class Manager {
   public proxyMap = new Map<string, ProxyEntry>();
   public componentMap = new Map<string, ComponentEntry>();
+  public setup = () => {};
 
   constructor(
-    public options: {
-      hot: ViteHot;
-      runtime: Runtime;
-      debug?: boolean;
-    }
+    public runtime: Runtime,
+    public options: RefreshRuntimeOptions
   ) {}
 
   wrap(name: string, Component: FC, key: string): FC {
@@ -49,16 +58,6 @@ export class Manager {
     return proxy.Component;
   }
 
-  setup() {
-    // https://vitejs.dev/guide/api-hmr.html#hot-accept-cb
-    this.options.hot.accept((newModule) => {
-      const ok = newModule && this.patch();
-      if (!ok) {
-        this.options.hot.invalidate();
-      }
-    });
-  }
-
   patch() {
     // TODO: debounce re-rendering?
     const componentNames = new Set([
@@ -84,16 +83,8 @@ export class Manager {
   }
 }
 
-export function createManager(
-  hot: ViteHot,
-  runtime: Runtime,
-  debug?: boolean
-): Manager {
-  return (hot.data[MANAGER_KEY] ??= new Manager({ hot, runtime, debug }));
-}
-
 function createProxyComponent(manager: Manager, name: string): ProxyEntry {
-  const { createElement, useEffect, useReducer } = manager.options.runtime;
+  const { createElement, useEffect, useReducer } = manager.runtime;
 
   const listeners = new Set<() => void>();
 
@@ -128,3 +119,63 @@ function createProxyComponent(manager: Manager, name: string): ProxyEntry {
 
   return { Component, listeners };
 }
+
+//
+// HMR API integration
+//
+
+export function initialize(
+  hot: ViteHot | WebpackHot,
+  runtime: Runtime,
+  options: RefreshRuntimeOptions
+) {
+  if (options.mode === "vite") {
+    return initializeVite(hot as any, runtime, options);
+  }
+  if (options.mode === "webpack") {
+    return initializeWebpack(hot as any, runtime, options);
+  }
+  return options.mode satisfies never;
+}
+
+function initializeVite(
+  hot: ViteHot,
+  runtime: Runtime,
+  options: RefreshRuntimeOptions
+) {
+  const manager = (hot.data[MANAGER_KEY] ??= new Manager(runtime, options));
+
+  // https://vitejs.dev/guide/api-hmr.html#hot-accept-cb
+  hot.accept((newModule) => {
+    const ok = newModule && manager.patch();
+    if (!ok) {
+      hot.invalidate();
+    }
+  });
+
+  return manager;
+}
+
+function initializeWebpack(
+  hot: WebpackHot,
+  runtime: Runtime,
+  options: RefreshRuntimeOptions
+) {
+  // 'hot.data' is passed from old module via hot.dispose(data)
+  // https://webpack.js.org/api/hot-module-replacement/#dispose-or-adddisposehandler
+  const prevData = hot.data?.[MANAGER_KEY];
+  const manager = prevData ?? new Manager(runtime, options);
+
+  hot.accept();
+  hot.dispose((data) => {
+    data[MANAGER_KEY] = manager;
+  });
+
+  manager.setup = () => {
+    if (prevData && !manager.patch()) {
+      hot.invalidate();
+    }
+  };
+
+  return manager;
+}

packages/tiny-refresh/src/transform.test.ts

@@ -1,7 +1,7 @@
 import { describe, expect, it } from "vitest";
-import { hmrTransform } from "./transform";
+import { transform } from "./transform";
 
-describe(hmrTransform, () => {
+describe(transform, () => {
   it("basic", async () => {
     const input = /* js */ `\
 
@@ -28,9 +28,11 @@ const NotFn = "hello";
 // export const NotFn2 = "hello";
 `;
     expect(
-      await hmrTransform(input, {
+      await transform(input, {
         runtime: "/runtime",
         refreshRuntime: "/refresh-runtime",
+        mode: "vite",
+        debug: false,
       })
     ).toMatchInlineSnapshot(`
       "
@@ -59,20 +61,18 @@ const NotFn = "hello";
       import * as $$runtime from "/runtime";
       import * as $$refresh from "/refresh-runtime";
       if (import.meta.hot) {
-        () => import.meta.hot.accept();
-        const $$manager = $$refresh.createManager(
+        (() => import.meta.hot.accept());
+        const $$manager = $$refresh.initialize(
           import.meta.hot,
-          {
-            createElement: $$runtime.createElement,
-            useReducer: $$runtime.useReducer,
-            useEffect: $$runtime.useEffect,
-          },
-          false,
+          $$runtime,
+          {"runtime":"/runtime","refreshRuntime":"/refresh-runtime","mode":"vite","debug":false}
         );
+
         FnDefault = $$manager.wrap("FnDefault", FnDefault, "");
         FnLet = $$manager.wrap("FnLet", FnLet, "useState/useRef/useCallback");
         FnConst = $$manager.wrap("FnConst", FnConst, "");
         FnNonExport = $$manager.wrap("FnNonExport", FnNonExport, "");
+
         $$manager.setup();
       }
       "