fix(hmr): support multiple accept calls (vitejs#170)

fix vitejs#158

playground/testHmrManual.js

@@ -1,13 +1,26 @@
 import { hot } from 'vite/hmr'
+import './testHmrManualDep'
 
 export const foo = 1
 
 if (__DEV__) {
   hot.accept(({ foo }) => {
-    console.log('foo is now: ', foo)
+    console.log('(self-accepting)1.foo is now:', foo)
+  })
+
+  hot.accept(({ foo }) => {
+    console.log('(self-accepting)2.foo is now:', foo)
   })
 
   hot.dispose(() => {
-    console.log('foo was: ', foo)
+    console.log(`foo was: ${foo}`)
+  })
+
+  hot.accept('./testHmrManualDep.js', ({ foo }) => {
+    console.log('(single dep) foo is now:', foo)
+  })
+
+  hot.accept(['./testHmrManualDep.js'], (modules) => {
+    console.log('(multiple deps) foo is now:', modules[0].foo)
   })
 }

playground/testHmrManualDep.js

@@ -0,0 +1,9 @@
+import { hot } from 'vite/hmr'
+
+export const foo = 1
+
+if (__DEV__) {
+  hot.dispose(() => {
+    console.log(`(dep) foo was: ${foo}`)
+  })
+}

src/client/client.ts

@@ -119,15 +119,7 @@ socket.addEventListener('message', async ({ data }) => {
       }
       break
     case 'js-update':
-      const update = jsUpdateMap.get(path)
-      if (update) {
-        update(timestamp)
-        console.log(`[vite]: js module hot updated: `, path)
-      } else {
-        console.error(
-          `[vite] got js update notification but no client callback was registered. Something is wrong.`
-        )
-      }
+      await updateModule(path, changeSrcPath, timestamp)
       break
     case 'custom':
       const cbs = customUpdateMap.get(id)
@@ -166,33 +158,97 @@ export function updateStyle(id: string, url: string) {
   link.setAttribute('href', url)
 }
 
-const jsUpdateMap = new Map<string, (timestamp: number) => void>()
+async function updateModule(
+  id: string,
+  changedPath: string,
+  timestamp: string
+) {
+  const mod = jsHotModuleMap.get(id)
+  if (!mod) {
+    console.error(
+      `[vite] got js update notification but no client callback was registered. Something is wrong.`
+    )
+    return
+  }
+
+  const moduleMap = new Map()
+  const isSelfUpdate = id === changedPath
+
+  // make sure we only import each dep once
+  const modulesToUpdate = new Set<string>()
+  if (isSelfUpdate) {
+    // self update - only update self
+    modulesToUpdate.add(id)
+  } else {
+    // dep update
+    for (const { deps } of mod.callbacks) {
+      if (Array.isArray(deps)) {
+        deps.forEach((dep) => modulesToUpdate.add(dep))
+      } else if (deps !== id) {
+        // exclude self accept calls
+        modulesToUpdate.add(deps)
+      }
+    }
+  }
+
+  // determine the qualified callbacks before we re-import the modules
+  const callbacks = mod.callbacks.filter(({ deps }) => {
+    return Array.isArray(deps)
+      ? deps.some((dep) => modulesToUpdate.has(dep))
+      : modulesToUpdate.has(deps)
+  })
+  // reset callbacks on self update since they are going to be registered again
+  if (isSelfUpdate) {
+    mod.callbacks = []
+  }
+
+  await Promise.all(
+    Array.from(modulesToUpdate).map(async (dep) => {
+      debugger
+      const disposer = jsDisposeMap.get(dep)
+      if (disposer) await disposer()
+      const newMod = await import(dep + `?t=${timestamp}`)
+      moduleMap.set(dep, newMod)
+    })
+  )
+
+  for (const { deps, fn } of callbacks) {
+    if (Array.isArray(deps)) {
+      fn(deps.map((dep) => moduleMap.get(dep)))
+    } else {
+      fn(moduleMap.get(deps))
+    }
+  }
+
+  console.log(`[vite]: js module hot updated: `, id)
+}
+
+interface HotModule {
+  id: string
+  callbacks: HotCallback[]
+}
+
+interface HotCallback {
+  deps: string | string[]
+  fn: (modules: object | object[]) => void
+}
+
+const jsHotModuleMap = new Map<string, HotModule>()
 const jsDisposeMap = new Map<string, () => void | Promise<void>>()
 const customUpdateMap = new Map<string, ((customData: any) => void)[]>()
 
 export const hot = {
   accept(
     id: string,
-    deps: string | string[],
-    callback: (modules: object | object[]) => void = () => {}
+    deps: HotCallback['deps'],
+    callback: HotCallback['fn'] = () => {}
   ) {
-    jsUpdateMap.set(id, async (timestamp: number) => {
-      if (Array.isArray(deps)) {
-        for (const id of deps) {
-          const disposer = jsDisposeMap.get(id)
-          if (disposer) await disposer()
-        }
-        Promise.all(deps.map((dep) => import(dep + `?t=${timestamp}`)))
-          .then(callback)
-          .catch((err) => warnFailedFetch(err, deps))
-      } else {
-        const disposer = jsDisposeMap.get(deps)
-        disposer && (await disposer())
-        import(deps + `?t=${timestamp}`)
-          .then(callback)
-          .catch((err) => warnFailedFetch(err, deps))
-      }
-    })
+    const mod: HotModule = jsHotModuleMap.get(id) || {
+      id,
+      callbacks: []
+    }
+    mod.callbacks.push({ deps, fn: callback })
+    jsHotModuleMap.set(id, mod)
   },
 
   dispose(id: string, cb: () => void) {

src/node/server/serverPluginHmr.ts

@@ -316,26 +316,26 @@ export const hmrPlugin: ServerPlugin = ({
         })
         console.log(chalk.green(`[vite] `) + `page reloaded.`)
       } else {
-        vueBoundaries.forEach((vueImporter) => {
+        vueBoundaries.forEach((vueBoundary) => {
           console.log(
             chalk.green(`[vite:hmr] `) +
-              `${vueImporter} reloaded due to change in ${relativeFile}.`
+              `${vueBoundary} reloaded due to change in ${relativeFile}.`
           )
           send({
             type: 'vue-reload',
-            path: vueImporter,
+            path: vueBoundary,
             changeSrcPath: publicPath,
             timestamp
           })
         })
-        jsBoundaries.forEach((jsImporter) => {
+        jsBoundaries.forEach((jsBoundary) => {
           console.log(
             chalk.green(`[vite:hmr] `) +
-              `${jsImporter} updated due to change in ${relativeFile}.`
+              `${jsBoundary} updated due to change in ${relativeFile}.`
           )
           send({
             type: 'js-update',
-            path: jsImporter,
+            path: jsBoundary,
             changeSrcPath: publicPath,
             timestamp
           })

test/test.js

@@ -133,17 +133,39 @@ describe('vite', () => {
         await expectByPolling(() => getText('.hmr-propagation'), '666')
       })
 
-      test('hmr (manual API)', async () => {
+      test('hmr (manual API, self accepting)', async () => {
         await updateFile('testHmrManual.js', (content) =>
           content.replace('foo = 1', 'foo = 2')
         )
         await expectByPolling(
           () => browserLogs[browserLogs.length - 1],
-          'foo is now:  2'
+          'js module hot updated:  /testHmrManual.js'
         )
-        // there will be a "js module reloaded" message in between because
-        // disposers are called before the new module is loaded.
-        expect(browserLogs[browserLogs.length - 3]).toMatch('foo was:  1')
+        expect(
+          browserLogs.slice(browserLogs.length - 4, browserLogs.length - 1)
+        ).toEqual([
+          `foo was: 1`,
+          `(self-accepting)1.foo is now: 2`,
+          `(self-accepting)2.foo is now: 2`
+        ])
+      })
+
+      test('hmr (manual API, accepting deps)', async () => {
+        browserLogs.length = 0
+        await updateFile('testHmrManualDep.js', (content) =>
+          content.replace('foo = 1', 'foo = 2')
+        )
+        await expectByPolling(
+          () => browserLogs[browserLogs.length - 1],
+          'js module hot updated:  /testHmrManual.js'
+        )
+        expect(
+          browserLogs.slice(browserLogs.length - 4, browserLogs.length - 1)
+        ).toEqual([
+          `(dep) foo was: 1`,
+          `(single dep) foo is now: 2`,
+          `(multiple deps) foo is now: 2`
+        ])
       })
     }
 
@@ -414,7 +436,7 @@ async function updateFile(file, replacer) {
 async function expectByPolling(poll, expected) {
   const maxTries = 20
   for (let tries = 0; tries < maxTries; tries++) {
-    const actual = await poll()
+    const actual = (await poll()) || ''
     if (actual.indexOf(expected) > -1 || tries === maxTries - 1) {
       expect(actual).toMatch(expected)
       break