From c92b55ecff1f31e00522c41bd59b02830ae74eb8 Mon Sep 17 00:00:00 2001
From: Anthony Fu <anthonyfu117@hotmail.com>
Date: Thu, 18 Aug 2022 15:25:51 +0800
Subject: [PATCH 1/2] feat(vite-node): record importers

---
 packages/vite-node/src/client.ts | 38 ++++++++++++++++++++------------
 packages/vite-node/src/types.ts  |  4 ++++
 2 files changed, 28 insertions(+), 14 deletions(-)

diff --git a/packages/vite-node/src/client.ts b/packages/vite-node/src/client.ts
index de18e6514529..6e65491a9c33 100644
--- a/packages/vite-node/src/client.ts
+++ b/packages/vite-node/src/client.ts
@@ -63,9 +63,11 @@ export class ModuleCacheMap extends Map<string, ModuleCache> {
     return super.set(fsPath, mod)
   }
 
-  get(fsPath: string) {
+  get(fsPath: string): ModuleCache {
     fsPath = this.normalizePath(fsPath)
-    return super.get(fsPath)
+    if (!super.has(fsPath))
+      super.set(fsPath, {})
+    return super.get(fsPath)!
   }
 
   delete(fsPath: string) {
@@ -104,16 +106,24 @@ export class ViteNodeRunner {
     const id = normalizeRequestId(rawId, this.options.base)
     const fsPath = toFilePath(id, this.root)
 
+    const mod = this.moduleCache.get(fsPath)
+    const importee = callstack[callstack.length - 1]
+
+    if (!mod.importers)
+      mod.importers = new Set()
+    if (importee)
+      mod.importers.add(importee)
+
     // the callstack reference itself circularly
-    if (callstack.includes(fsPath) && this.moduleCache.get(fsPath)?.exports)
-      return this.moduleCache.get(fsPath)?.exports
+    if (callstack.includes(fsPath) && mod.exports)
+      return mod.exports
 
     // cached module
-    if (this.moduleCache.get(fsPath)?.promise)
-      return this.moduleCache.get(fsPath)?.promise
+    if (mod.promise)
+      return mod.promise
 
     const promise = this.directRequest(id, fsPath, callstack)
-    this.moduleCache.update(fsPath, { promise })
+    Object.assign(mod, { promise })
 
     return await promise
   }
@@ -122,6 +132,8 @@ export class ViteNodeRunner {
   async directRequest(id: string, fsPath: string, _callstack: string[]) {
     const callstack = [..._callstack, fsPath]
 
+    const mod = this.moduleCache.get(fsPath)
+
     const request = async (dep: string) => {
       const depFsPath = toFilePath(normalizeRequestId(dep, this.options.base), this.root)
       const getStack = () => {
@@ -140,9 +152,7 @@ export class ViteNodeRunner {
           throw new Error(`[vite-node] Failed to resolve circular dependency, ${getStack()}`)
         }
 
-        const mod = await this.cachedRequest(dep, callstack)
-
-        return mod
+        return await this.cachedRequest(dep, callstack)
       }
       finally {
         if (debugTimer)
@@ -179,9 +189,9 @@ export class ViteNodeRunner {
     let { code: transformed, externalize } = await this.options.fetchModule(id)
     if (externalize) {
       debugNative(externalize)
-      const mod = await this.interopedImport(externalize)
-      this.moduleCache.update(fsPath, { exports: mod })
-      return mod
+      const exports = await this.interopedImport(externalize)
+      mod.exports = exports
+      return exports
     }
 
     if (transformed == null)
@@ -197,7 +207,7 @@ export class ViteNodeRunner {
       configurable: false,
     })
 
-    this.moduleCache.update(fsPath, { code: transformed, exports })
+    Object.assign(mod, { code: transformed, exports })
 
     const __filename = fileURLToPath(url)
     const moduleProxy = {
diff --git a/packages/vite-node/src/types.ts b/packages/vite-node/src/types.ts
index 030808e59f0a..b6581a641786 100644
--- a/packages/vite-node/src/types.ts
+++ b/packages/vite-node/src/types.ts
@@ -45,6 +45,10 @@ export interface ModuleCache {
   promise?: Promise<any>
   exports?: any
   code?: string
+  /**
+   * Module ids that imports this module
+   */
+  importers?: Set<string>
 }
 
 export interface ViteNodeRunnerOptions {

From 3c2907911e7bd609b7b1683d2d192ff410fb7c65 Mon Sep 17 00:00:00 2001
From: Anthony Fu <anthonyfu117@hotmail.com>
Date: Thu, 18 Aug 2022 15:44:08 +0800
Subject: [PATCH 2/2] feat: invalidateDepTree

---
 packages/vite-node/src/client.ts | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/packages/vite-node/src/client.ts b/packages/vite-node/src/client.ts
index 6e65491a9c33..822dc189ac84 100644
--- a/packages/vite-node/src/client.ts
+++ b/packages/vite-node/src/client.ts
@@ -74,6 +74,23 @@ export class ModuleCacheMap extends Map<string, ModuleCache> {
     fsPath = this.normalizePath(fsPath)
     return super.delete(fsPath)
   }
+
+  /**
+   * Invalidate modules that dependent on the given modules, up to the main entry
+   */
+  invalidateDepTree(ids: string[] | Set<string>, invalidated = new Set<string>()) {
+    for (const _id of ids) {
+      const id = this.normalizePath(_id)
+      if (invalidated.has(id))
+        continue
+      invalidated.add(id)
+      const mod = super.get(id)
+      if (mod?.importers)
+        this.invalidateDepTree(mod.importers, invalidated)
+      super.delete(id)
+    }
+    return invalidated
+  }
 }
 
 export class ViteNodeRunner {