feat: support self-accepting hmr

README.md

@@ -78,13 +78,32 @@ The above will throw an error by default. `vite` detects such bare module import
 
   foo()
 
-  hot.accept('./foo.js', ({ foo }) => {
+  hot.accept('./foo.js', (newFoo) => {
     // the callback receives the updated './foo.js' module
-    foo()
+    newFoo.foo()
+  })
+
+  // Can also accept an array of dep modules:
+  hot.accept(['./foo.js', './bar.js'], ([newFooModule, newBarModule]) => {
+    // the callback receives the updated mdoules in an Array
   })
   ```
 
-  Note it's simplified and not fully compatible with webpack's HMR API, for example there is no self-accepting modules, and if you re-export `foo` from this file, it won't reflect changes in modules that import this file.
+  Modules can also be self-accepting:
+
+  ```js
+  import { hot } from '@hmr'
+
+  export const count = 1
+
+  hot.accept(newModule => {
+    console.log('updated: count is now ', newModule.count)
+  })
+  ```
+
+  Note that `vite`'s HMR does not actually swap the originally imported module: if an accepting module re-exports imports from a dep, then it is responsible for updating those re-exports (and these exports must be using `let`). In addition, importers up the chain from the accepting module will not be notified of the change.
+
+  This simplified HMR implementation is sufficient for most dev use cases, while allowing us to skip the expensive work of generating proxy modules.
 
 ### CSS Pre-Processors
 

src/client/client.ts

@@ -90,7 +90,7 @@ export const hot = {
   accept(
     importer: string,
     deps: string | string[],
-    callback: (modules: object | object[]) => void
+    callback: (modules: object | object[]) => void = () => {}
   ) {
     jsUpdateMap.set(importer, (timestamp: number) => {
       if (Array.isArray(deps)) {

src/node/serverPluginHmr.ts

@@ -54,7 +54,7 @@ export type HMRWatcher = FSWatcher & {
 // so that we can determine what files to hot reload
 type HMRStateMap = Map<string, Set<string>>
 
-export const hmrBoundariesMap: HMRStateMap = new Map()
+export const hmrAcceptanceMap: HMRStateMap = new Map()
 export const importerMap: HMRStateMap = new Map()
 export const importeeMap: HMRStateMap = new Map()
 
@@ -260,11 +260,17 @@ function walkImportChain(
   vueImporters: Set<string>,
   jsHotImporters: Set<string>
 ): boolean {
+  if (isHmrAccepted(importee, importee)) {
+    // self-accepting module.
+    jsHotImporters.add(importee)
+    return false
+  }
+
   let hasDeadEnd = false
   for (const importer of currentImporters) {
     if (importer.endsWith('.vue')) {
       vueImporters.add(importer)
-    } else if (isHMRBoundary(importer, importee)) {
+    } else if (isHmrAccepted(importer, importee)) {
       jsHotImporters.add(importer)
     } else {
       const parentImpoters = importerMap.get(importer)
@@ -283,8 +289,8 @@ function walkImportChain(
   return hasDeadEnd
 }
 
-function isHMRBoundary(importer: string, dep: string): boolean {
-  const deps = hmrBoundariesMap.get(importer)
+function isHmrAccepted(importer: string, dep: string): boolean {
+  const deps = hmrAcceptanceMap.get(importer)
   return deps ? deps.has(dep) : false
 }
 
@@ -327,7 +333,7 @@ export function rewriteFileWithHMR(
   }).program.body
 
   const registerDep = (e: StringLiteral) => {
-    const deps = ensureMapEntry(hmrBoundariesMap, importer)
+    const deps = ensureMapEntry(hmrAcceptanceMap, importer)
     const depPublicPath = slash(path.resolve(path.dirname(importer), e.value))
     deps.add(depPublicPath)
     debugHmr(`        ${importer} accepts ${depPublicPath}`)
@@ -359,9 +365,15 @@ export function rewriteFileWithHMR(
         })
       } else if (args[0].type === 'StringLiteral') {
         registerDep(args[0])
+      } else if (args[0].type.endsWith('FunctionExpression')) {
+        // self accepting, rewrite to inject itself
+        // hot.accept(() => {})  -->  hot.accept('/foo.js', '/foo.js', () => {})
+        s.appendLeft(args[0].start!, JSON.stringify(importer) + ', ')
+        ensureMapEntry(hmrAcceptanceMap, importer).add(importer)
       } else {
         console.error(
-          `[vite] HMR syntax error in ${importer}: hot.accept() expects a dep string or an array of deps.`
+          `[vite] HMR syntax error in ${importer}: ` +
+            `hot.accept() expects a dep string, an array of deps, or a callback.`
         )
       }
     }