From 5913ebc34f810e84b4ad482aa835e1a4e8beb404 Mon Sep 17 00:00:00 2001
From: Evan You <yyx990803@gmail.com>
Date: Tue, 28 Feb 2023 07:56:27 +0800
Subject: [PATCH] fix: hmr on deps change of data loaders

---
 src/node/plugins/staticDataPlugin.ts | 37 ++++++++++++++++++++--------
 1 file changed, 27 insertions(+), 10 deletions(-)

diff --git a/src/node/plugins/staticDataPlugin.ts b/src/node/plugins/staticDataPlugin.ts
index bf0261bb..adf86722 100644
--- a/src/node/plugins/staticDataPlugin.ts
+++ b/src/node/plugins/staticDataPlugin.ts
@@ -4,7 +4,7 @@ import {
   loadConfigFromFile,
   normalizePath
 } from 'vite'
-import { dirname, resolve } from 'path'
+import path, { dirname, resolve } from 'path'
 import { isMatch } from 'micromatch'
 
 const loaderMatch = /\.data\.(j|t)s$/
@@ -24,6 +24,8 @@ interface CachedLoaderModule {
 const idToLoaderModulesMap: Record<string, CachedLoaderModule | undefined> =
   Object.create(null)
 
+const depToLoaderModuleIdMap: Record<string, string> = Object.create(null)
+
 // During build, the load hook will be called on the same file twice
 // once for client and once for server build. Not only is this wasteful, it
 // also leads to a race condition in loadConfigFromFile() that results in an
@@ -66,8 +68,16 @@ export const staticDataPlugin: Plugin = {
       } else {
         // use vite's load config util as a away to load Node.js file with
         // TS & native ESM support
-        const loaderModule = (await loadConfigFromFile({} as any, id))
-          ?.config as LoaderModule
+        const res = await loadConfigFromFile({} as any, id)
+
+        // record deps for hmr
+        if (res) {
+          for (const dep of res.dependencies) {
+            depToLoaderModuleIdMap[normalizePath(path.resolve(dep))] = id
+          }
+        }
+
+        const loaderModule = res?.config as LoaderModule
         pattern =
           typeof loaderModule.watch === 'string'
             ? [loaderModule.watch]
@@ -103,20 +113,27 @@ export const staticDataPlugin: Plugin = {
     if (server && loaderMatch.test(id)) {
       // register this module as a glob importer
       const { pattern } = idToLoaderModulesMap[id]!
-      ;(server as any)._importGlobMap.set(id, [pattern])
+      if (pattern) {
+        ;(server as any)._importGlobMap.set(id, [pattern])
+      }
     }
     return null
   },
 
   handleHotUpdate(ctx) {
+    const file = normalizePath(ctx.file)
+
+    // dependency of data loader changed
+    // (note the dep array includes the loader file itself)
+    if (file in depToLoaderModuleIdMap) {
+      const id = depToLoaderModuleIdMap[file]!
+      delete idToLoaderModulesMap[id]
+      ctx.modules.push(server.moduleGraph.getModuleById(id)!)
+    }
+
     for (const id in idToLoaderModulesMap) {
       const { pattern } = idToLoaderModulesMap[id]!
-      const isLoaderFile = normalizePath(ctx.file) === id
-      if (isLoaderFile) {
-        // invalidate loader file
-        delete idToLoaderModulesMap[id]
-      }
-      if (isLoaderFile || (pattern && isMatch(ctx.file, pattern))) {
+      if (pattern && isMatch(file, pattern)) {
         ctx.modules.push(server.moduleGraph.getModuleById(id)!)
       }
     }