From 50f0700d241a5d5b443e38f3ef7ce50640a993fa Mon Sep 17 00:00:00 2001
From: Dunqing <dengqing0821@gmail.com>
Date: Tue, 6 Jun 2023 23:44:31 +0800
Subject: [PATCH] fix(vite-node): circular import stuck (#3480)

---
 examples/mocks/src/main.js           |  5 +++++
 examples/mocks/test/circular.spec.ts | 11 +++++++++-
 packages/vite-node/src/client.ts     | 32 +++++++++++++++++++---------
 packages/vite-node/src/types.ts      |  1 +
 4 files changed, 38 insertions(+), 11 deletions(-)

diff --git a/examples/mocks/src/main.js b/examples/mocks/src/main.js
index be77a608ffc2..5401d983c887 100644
--- a/examples/mocks/src/main.js
+++ b/examples/mocks/src/main.js
@@ -1,5 +1,10 @@
 import { funcA } from './A'
+import { funcB } from './B'
 
 export function main() {
   return funcA()
 }
+
+export function mainB() {
+  return funcB()
+}
diff --git a/examples/mocks/test/circular.spec.ts b/examples/mocks/test/circular.spec.ts
index 3c34b0916105..e8f4651dfef9 100644
--- a/examples/mocks/test/circular.spec.ts
+++ b/examples/mocks/test/circular.spec.ts
@@ -1,5 +1,5 @@
 import { expect, test, vi } from 'vitest'
-import { main } from '../src/main.js'
+import { main, mainB } from '../src/main.js'
 import x from '../src/export-default-circle-index'
 
 vi.mock('../src/A', async () => ({
@@ -7,12 +7,21 @@ vi.mock('../src/A', async () => ({
   funcA: () => 'mockedA',
 }))
 
+vi.mock('../src/B', async () => ({
+  ...(await vi.importActual<any>('../src/B')),
+  funcB: () => 'mockedB',
+}))
+
 vi.mock('../src/export-default-circle-b')
 
 test('can import actual inside mock factory', () => {
   expect(main()).toBe('mockedA')
 })
 
+test('can import in top level and inside mock factory', () => {
+  expect(mainB()).toBe('mockedB')
+})
+
 test('can mock a circular dependency', () => {
   expect(x()).toBe(undefined)
 })
diff --git a/packages/vite-node/src/client.ts b/packages/vite-node/src/client.ts
index 3f2fbc139791..c6ba7ebc0da7 100644
--- a/packages/vite-node/src/client.ts
+++ b/packages/vite-node/src/client.ts
@@ -58,10 +58,10 @@ export class ModuleCacheMap extends Map<string, ModuleCache> {
   /**
    * Assign partial data to the map
    */
-  update(fsPath: string, mod: Partial<ModuleCache>) {
+  update(fsPath: string, mod: ModuleCache) {
     fsPath = this.normalizePath(fsPath)
     if (!super.has(fsPath))
-      super.set(fsPath, mod)
+      this.setByModuleId(fsPath, mod)
     else
       Object.assign(super.get(fsPath) as ModuleCache, mod)
     return this
@@ -75,13 +75,21 @@ export class ModuleCacheMap extends Map<string, ModuleCache> {
     return this.setByModuleId(this.normalizePath(fsPath), mod)
   }
 
-  getByModuleId(modulePath: string): ModuleCache {
+  getByModuleId(modulePath: string) {
     if (!super.has(modulePath))
-      super.set(modulePath, {})
-    return super.get(modulePath)!
+      this.setByModuleId(modulePath, {})
+
+    const mod = super.get(modulePath)!
+    if (!mod.imports) {
+      Object.assign(mod, {
+        imports: new Set(),
+        importers: new Set(),
+      })
+    }
+    return mod as ModuleCache & Required<Pick<ModuleCache, 'imports' | 'importers'>>
   }
 
-  get(fsPath: string): ModuleCache {
+  get(fsPath: string) {
     return this.getByModuleId(this.normalizePath(fsPath))
   }
 
@@ -99,6 +107,7 @@ export class ModuleCacheMap extends Map<string, ModuleCache> {
     delete mod.promise
     delete mod.exports
     mod.importers?.clear()
+    mod.imports?.clear()
     return true
   }
 
@@ -185,16 +194,15 @@ export class ViteNodeRunner {
     const importee = callstack[callstack.length - 1]
 
     const mod = this.moduleCache.get(fsPath)
+    const { imports, importers } = mod
 
-    if (!mod.importers)
-      mod.importers = new Set()
     if (importee)
-      mod.importers.add(importee)
+      importers.add(importee)
 
     const getStack = () => `stack:\n${[...callstack, fsPath].reverse().map(p => `  - ${p}`).join('\n')}`
 
     // check circular dependency
-    if (callstack.includes(fsPath) || callstack.some(c => this.moduleCache.get(c).importers?.has(fsPath))) {
+    if (callstack.includes(fsPath) || Array.from(imports.values()).some(i => importers.has(i))) {
       if (mod.exports)
         return mod.exports
     }
@@ -269,6 +277,10 @@ export class ViteNodeRunner {
 
     const request = async (dep: string) => {
       const [id, depFsPath] = await this.resolveUrl(`${dep}`, fsPath)
+      const depMod = this.moduleCache.getByModuleId(depFsPath)
+      depMod.importers.add(moduleId)
+      mod.imports.add(depFsPath)
+
       return this.dependencyRequest(id, depFsPath, callstack)
     }
 
diff --git a/packages/vite-node/src/types.ts b/packages/vite-node/src/types.ts
index 077f974ec0a0..18d6016a5377 100644
--- a/packages/vite-node/src/types.ts
+++ b/packages/vite-node/src/types.ts
@@ -66,6 +66,7 @@ export interface ModuleCache {
    * Module ids that imports this module
    */
   importers?: Set<string>
+  imports?: Set<string>
 }
 
 export interface ViteNodeRunnerOptions {